Use irange in a few places (#55325)

r-barnes · facebook-github-bot · commit 93140a31e25c · 2021-06-07T14:53:41.000-07:00
Summary: Pull Request resolved: #55325 Test Plan: Sandcastle Reviewed By: SciPioneer Differential Revision: D27573006 fbshipit-source-id: 647b5da3901e92c23e95b2fe5e833e9081d72837
diff --git a/torch/lib/c10d/ProcessGroupMPI.cpp b/torch/lib/c10d/ProcessGroupMPI.cpp
@@ -4,6 +4,7 @@
 #include <map>
 
 #include <c10/core/DeviceGuard.h>
+#include <c10/util/irange.h>
 
 #if defined(OPEN_MPI) && OPEN_MPI
 #include <mpi-ext.h> // Needed for CUDA-aware check
@@ -783,7 +784,7 @@ c10::intrusive_ptr<ProcessGroup::Work> ProcessGroupMPI::alltoall(
         at::Tensor dstFlatData = at::empty({dst_len}, dstdata[0].options());
         auto srcFlatDataSplits =
             srcFlatData.split_with_sizes(c10::IntArrayRef(send_lengthsL), 0);
-        for (int i = 0; i < size_; i++) {
+        for(const auto i : c10::irange(size_)) {
           srcFlatDataSplits[i].copy_(srcdata[i].view({-1}));
         }
         c10::DeviceGuard guard1(srcdata[0].device());
@@ -801,7 +802,7 @@ c10::intrusive_ptr<ProcessGroup::Work> ProcessGroupMPI::alltoall(
 
         auto dstFlatDataSplits =
             dstFlatData.split_with_sizes(c10::IntArrayRef(recv_lengthsL), 0);
-        for (int i = 0; i < size_; i++) {
+        for(const auto i : c10::irange(size_)) {
           dstdata[i].view({-1}).copy_(dstFlatDataSplits[i]);
         }
       };
diff --git a/torch/lib/c10d/Utils.hpp b/torch/lib/c10d/Utils.hpp
@@ -2,6 +2,7 @@
 
 #include <ATen/ATen.h>
 #include <c10/util/accumulate.h>
+#include <c10/util/irange.h>
 #include <c10d/Types.hpp>
 
 #ifdef _WIN32
@@ -447,7 +448,7 @@ size_t computeLengthsAndOffsets(
     equal_splits = true;
     split_size = tensor.size(0) / group_size;
   }
-  for (int i = 0; i < group_size; i++) {
+  for(const auto i : c10::irange(group_size)) {
     size_t length = row_size * (equal_splits ? split_size : split_sizes[i]);
     TORCH_INTERNAL_ASSERT(
         length <= std::numeric_limits<int>::max() &&
@@ -467,7 +468,7 @@ size_t computeLengthsAndOffsets(
     std::vector<T>* offsets) {
   size_t group_size = lengths->size();
   size_t offset = 0;
-  for (int i = 0; i < group_size; i++) {
+  for(const auto i : c10::irange(group_size)) {
     size_t length = tensors[i].numel();
     TORCH_INTERNAL_ASSERT(
         length <= std::numeric_limits<int>::max() &&
diff --git a/torch/lib/c10d/test/ProcessGroupGlooTest.cpp b/torch/lib/c10d/test/ProcessGroupGlooTest.cpp
@@ -16,6 +16,7 @@
 #include <torch/csrc/autograd/profiler.h>
 #include <torch/cuda.h>
 
+#include <c10/util/irange.h>
 #include <c10d/FileStore.hpp>
 #include <c10d/ProcessGroupGloo.hpp>
 #include <c10d/test/TestUtils.hpp>
@@ -509,7 +510,7 @@ void testMonitoredBarrier(const std::string& path) {
   };
   std::vector<std::thread> threads;
   threads.reserve(size);
-  for (int r = 0; r < size; r++) {
+  for(const auto r : c10::irange(size)) {
     threads.emplace_back(std::thread([=]() { runMonitoredBarrier(r); }));
   }
   for (auto & t : threads) {
@@ -530,7 +531,7 @@ void testMonitoredBarrier(const std::string& path) {
       }
   };
   threads.clear();
-  for (int r = 0; r < size; r++) {
+  for(const auto r : c10::irange(size)) {
       threads.emplace_back(std::thread([=]() { runMonitoredBarrierWithException(r); }));
   }
   for (auto & t : threads) {
diff --git a/torch/utils/benchmark/utils/timeit_template.cpp b/torch/utils/benchmark/utils/timeit_template.cpp
@@ -9,6 +9,7 @@ sections with user provided statements.
 */
 #include <chrono>
 
+#include <c10/util/irange.h>
 #include <pybind11/pybind11.h>
 #include <torch/extension.h>
 
@@ -28,7 +29,7 @@ double timeit(int n) {
 
     // Main loop
     auto start_time = std::chrono::high_resolution_clock::now();
-    for (int loop_idx = 0; loop_idx < n; loop_idx++) {
+    for(const auto loop_idx : c10::irange(n)) {
         // STMT_TEMPLATE_LOCATION
     }
     auto end_time = std::chrono::high_resolution_clock::now();
diff --git a/torch/utils/benchmark/utils/valgrind_wrapper/timer_callgrind_template.cpp b/torch/utils/benchmark/utils/valgrind_wrapper/timer_callgrind_template.cpp
@@ -8,11 +8,12 @@ This template will be consumed by `cpp_jit.py`, and will replace:
 sections with user provided statements.
 */
 
-#include <string>
-
 #include <callgrind.h>
+#include <c10/util/irange.h>
 #include <torch/torch.h>
 
+#include <string>
+
 // Global setup. (e.g. #includes)
 // GLOBAL_SETUP_TEMPLATE_LOCATION
 
@@ -42,15 +43,15 @@ int main(int argc, char* argv[]) {
     // SETUP_TEMPLATE_LOCATION
 
     // Warmup
-    for (int i = 0; i < number_warmup; i++) {
+    for(const auto i : c10::irange(number_warmup)) {
         // STMT_TEMPLATE_LOCATION
     }
 
     // Main loop
-    for (int repeat = 0; repeat < repeats; repeat++) {
+    for(const auto repeat : c10::irange(repeats)) {
         CALLGRIND_TOGGLE_COLLECT;
 
-        for (int i = 0; i < number; i++) {
+        for(const auto i : c10::irange(number)) {
         // STMT_TEMPLATE_LOCATION
         }
 
