[Inductor] Avoid tensor slice overflow for large step (pytorch#147433)

Fixes pytorch#147071

Currently, if step is a value very close to INT64_MAX, the calculation of slice output length will overflow. This PR tries to fix this problem and thus fix pytorch#147071.

Pull Request resolved: pytorch#147433
Approved by: https://github.com/leslie-fang-intel, https://github.com/jansel

Author: DDEle

aten/src/ATen/native/TensorShape.cpp

@@ -3059,7 +3059,8 @@ Tensor slice(
   }
   auto storage_offset = self.storage_offset() + start_val * strides[dim];
   auto len = end_val - start_val;
-  sizes[dim] = (len + step - 1) / step; // round-up
+  sizes[dim] =
+      (len == 0) ? 0 : (1 + (len - 1) / step); // round-up, avoiding overflow
   strides[dim] *= step;
 
   Tensor result;

test/expect/HasDecompTest.test_has_decomposition.expect

@@ -1167,8 +1167,6 @@ aten::set_
 aten::set_.source_Storage
 aten::set_.source_Storage_storage_offset
 aten::set_.source_Tensor
-aten::slice_copy.Tensor
-aten::slice_copy.Tensor_out
 aten::slice_inverse
 aten::slow_conv3d_forward
 aten::slow_conv3d_forward.output

test/inductor/test_torchinductor.py

@@ -12527,6 +12527,17 @@ def f(x):
             ms = do_bench(lambda: opt_f(x))
             print(f"{ms=:.3f}")
 
+    def test_slice_overflow(self):
+        # https://github.com/pytorch/pytorch/issues/147071
+        def f(input):
+            var = torch.slice_copy(
+                input, dim=0, start=449, end=None, step=9223372036854775807
+            )
+            return torch.reciprocal(var)
+
+        input = torch.randn((875,))
+        self.assertEqual(torch.compile(f)(input), f(input))
+
     @torch._inductor.config.patch("graph_partition", True)
     def test_graph_partition_no_inputs(self):
         def foo():

torch/_decomp/__init__.py

@@ -477,6 +477,7 @@ def _core_aten_decompositions_post_autograd() -> dict[
             aten.sinc,
             aten.sinc_,
             aten.slice_backward,
+            aten.slice_copy,
             aten.smooth_l1_loss,
             aten.smooth_l1_loss_backward,
             aten.soft_margin_loss,

torch/_decomp/decompositions.py

@@ -748,7 +748,7 @@ def slice_forward(
 
     storage_offset = self.storage_offset() + start_val * strides[dim]
     len = end_val - start_val
-    sizes[dim] = (len + step - 1) // step
+    sizes[dim] = -(len // -step)  # round-up, avoiding overflow
     strides[dim] *= step
 
     if self.is_quantized:
@@ -759,6 +759,23 @@ def slice_forward(
         return self.as_strided(sizes, strides, storage_offset)
 
 
+@register_decomposition([aten.slice_copy.Tensor, aten.slice_copy.Tensor_out])
+def slice_copy(
+    self: Tensor,
+    dim: int = 0,
+    start: Optional[int] = None,
+    end: Optional[int] = None,
+    step: int = 1,
+    out: Optional[Tensor] = None,
+):
+    _slice = slice_forward(self, dim, start, end, step)
+    slice_clone = _slice.clone(memory_format=torch.contiguous_format)
+    if out is None:
+        return slice_clone
+    else:
+        return _safe_copy_out(copy_from=slice_clone, copy_to=out, exact_dtype=True)
+
+
 def _normalize_start_end(
     x: Tensor, dim: int, start: Optional[int], end: Optional[int]
 ) -> tuple[int, int]: