update add ov:erflow check for stride calculation

pytorch · Apr 6, 2023 · f52b910 · f52b910
1 parent 66ac0a0
commit f52b910
Showing 1 changed file with 6 additions and 9 deletions.
diff --git a/c10/core/TensorImpl.h b/c10/core/TensorImpl.h
@@ -1754,7 +1754,7 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
         new_stride.size(),
         ")");
     const auto new_dim = new_size.size();
-
+    bool overflowed = false;
     sizes_and_strides_.set_sizes(new_size);
 
     if (new_dim > 0) {
@@ -1769,15 +1769,14 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
             sizes_and_strides_.stride_at_unchecked(dim) = 1;
           } else {
             // Keep stride monotonically increasing to match NumPy.
-            sizes_and_strides_.stride_at_unchecked(dim) =
-                std::max<int64_t>(
-                    sizes_and_strides_.size_at_unchecked(dim + 1), 1) *
-                sizes_and_strides_.stride_at_unchecked(dim + 1);
+            overflowed |= c10::mul_overflows(sizes_and_strides_.stride_at_unchecked(dim + 1), std::max<int64_t>(
+              sizes_and_strides_.size_at_unchecked(dim + 1), 1), std::addressof(sizes_and_strides_.stride_at_unchecked(dim)));
           }
         }
         if (dim == 0)
           break;
       }
+      TORCH_CHECK(!overflowed, "Stride calculation overflowed");
     }
 
     refresh_numel();
@@ -2279,11 +2278,9 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
           sizes_and_strides_.stride_at_unchecked(last_idx) = 1;
           for (auto i = last_idx - 1; i >= 0; --i) {
             overflowed |= c10::mul_overflows(sizes_and_strides_.stride_at_unchecked(i + 1), std::max<int64_t>(
-                    sizes_and_strides_.size_at_unchecked(i + 1), 1), std::addressof(sizes_and_strides_.stride_at_unchecked(i)));
-            TORCH_CHECK(!overflowed,
-              "Stride calculation overflowed with stride of dim", i);
+              sizes_and_strides_.size_at_unchecked(i + 1), 1), std::addressof(sizes_and_strides_.stride_at_unchecked(i)));
           }
-
+          TORCH_CHECK(!overflowed, "Stride calculation overflowed");
         }
         break;
       }