[acc_ops] Move slice_tensor to consider single dim at a time (#5906)

jfix71 · pytorchmergebot · commit 184b78c4c19d · 2022-01-31T23:37:36.000Z
Summary: Pull Request resolved: pytorch/glow#5906 Pull Request resolved: #71883 Fixes slice_tensor retracing. Include fix for retrace coverage. Missed in D33760455 (66939e3). Test Plan: CI Reviewed By: wushirong Differential Revision: D33802222 fbshipit-source-id: 4e0e44ae4a4eb70b99d79f0cd582182031b87e25 (cherry picked from commit 98fd23c)
diff --git a/test/fx_acc/test_acc_tracer.py b/test/fx_acc/test_acc_tracer.py
@@ -94,7 +94,7 @@ def forward(self, a: torch.Tensor) -> torch.Tensor:
 
         ref_outputs = m(a)
         outputs = traced(a)
-        traced_again = acc_tracer.trace(m, [a])
+        traced_again = acc_tracer.trace(traced, [a])
         outputs_again = traced_again(a)
         if isinstance(ref_outputs, torch.Tensor):
             ref_outputs = [ref_outputs]
@@ -1881,6 +1881,27 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         for i, j in zip(ref_output, output):
             self.assertTrue(torch.equal(i, j))
 
+    @parameterized.expand(
+        [
+            ("neg_1", -1, 1, 3),
+            ("neg_2", -2, 1, 3),
+            ("neg_4", -4, 1, 1),
+        ]
+    )
+    def test_negative_slicing(self, _, dim, start, length):
+        """
+        Test that slicing with negative dims works.
+        """
+        self._make_acc_op_function_test(
+            acc_ops.slice_tensor,
+            torch.narrow,
+            input_shape=(2, 3, 4, 5),
+            validate_same_kwargs=False,
+            dim=dim,
+            start=start,
+            length=length,
+        )
+
     def test_list_input(self):
         """
         Test that list inputs are traced correctly.
diff --git a/torch/fx/experimental/fx2trt/converters/acc_ops_converters.py b/torch/fx/experimental/fx2trt/converters/acc_ops_converters.py
@@ -1422,30 +1422,26 @@ def acc_ops_slice_tensor(
                            "of the TensorRT region!")
 
     ranks = len(input_val.shape) + (1 if network.has_implicit_batch_dimension else 0)
-    dims = [get_positive_dim(dim, ranks) for dim in cast(Sequence[int], kwargs["dims"])]
+    dim = get_positive_dim(cast(int, kwargs["dim"]), ranks)
 
     if network.has_implicit_batch_dimension:
-        if not len(dims):
-            raise RuntimeError("dim argument cannot be empty!")
-        if any([dim == 0 for dim in dims]):
+        if dim == 0:
             raise RuntimeError(
-                f"We do not support slice_tensor at batch dim when it's implicit, got {dims}!"
+                f"We do not support slice_tensor at batch dim when it's implicit, got {dim}!"
             )
-        dims = [d - 1 for d in dims]
+        dim = dim - 1
     else:
         raise RuntimeError("We don't support slice_tensor with explicit batch dimension yet!")
 
+    start_int = cast(int, kwargs["start"])
+    stop_int = cast(int, kwargs["stop"])
+    step_int = cast(int, kwargs["step"])
     start = [0] * len(input_val.shape)
+    start[dim] = start_int
     stride = [1] * len(start)
+    stride[dim] = step_int
     output_shape = list(input_val.shape)
-    starts = cast(Sequence[int], kwargs["starts"])
-    stops = cast(Sequence[int], kwargs["stops"])
-    steps = cast(Sequence[int], kwargs["steps"])
-
-    for i, dim in enumerate(dims):
-        start[dim] = starts[i]
-        stride[dim] = steps[i]
-        output_shape[dim] = (stops[i] - starts[i]) // steps[i]
+    output_shape[dim] = (stop_int - start_int) // step_int
 
     layer = network.add_slice(input_val, start=start, shape=output_shape, stride=stride)
     set_layer_name(layer, target, name)
diff --git a/torch/fx/experimental/fx_acc/acc_ops.py b/torch/fx/experimental/fx_acc/acc_ops.py
@@ -3,7 +3,7 @@
 import warnings
 
 import torch  # isort:skip
-from typing import Sequence, Optional, List, cast
+from typing import Sequence, List, cast
 
 import torch.fx.experimental.fx_acc.acc_utils as acc_utils
 import torch.nn as nn
@@ -1310,10 +1310,10 @@ def torch_split_mapper(node: torch.fx.Node, mod: nn.Module) -> torch.fx.Node:
             assert isinstance(i, int)
             new_kwargs = {
                 "input": node.kwargs["input"],
-                "dims": (node.kwargs["dim"],),
-                "starts": (start,),
-                "stops": (start + i,),
-                "steps": (1,),
+                "dim": node.kwargs["dim"],
+                "start": start,
+                "stop": start + i,
+                "step": 1,
             }
             new_node = node.graph.call_function(slice_tensor, kwargs=new_kwargs)
             new_node.meta["type"] = torch.Tensor
@@ -1504,19 +1504,16 @@ def getitem(*, input, idx):
 
 @register_acc_op_properties(AccOpProperty.unary)
 @register_acc_op
-def slice_tensor(*, input, dims, starts, stops, steps):
-    slices: List[Optional[slice]] = [None for _ in range(input.dim())]
-
-    # For all provided dims, extract out a slice for starts/stops/steps.
-    for idx, dim in enumerate(dims):
-        slices[dim] = slice(starts[idx], stops[idx], steps[idx])
-
-    # For all unspecified dims, default to the full slice.
-    for idx, s in enumerate(slices):
-        if s is None:
-            slices[idx] = slice(None, None, None)
+def slice_tensor(*, input, dim, start, stop, step):
+    slc = slice(start, stop, step)
+    if dim >= 0:
+        slices: List[slice] = [slice(None, None, None) for _ in range(dim)]
+        slices.append(slc)
+    else:
+        slices = [Ellipsis, slc]  # type: ignore[list-item]
+        slices.extend([slice(None, None, None) for _ in range(-dim - 1)])
 
-    return input[slices]
+    return input[tuple(slices)]
 
 
 @register_custom_acc_mapper_fn(
@@ -1543,10 +1540,10 @@ def custom_narrow_mapper(node: torch.fx.Node, mod: nn.Module) -> torch.fx.Node:
     )
     kwargs = {
         "input": node.kwargs["input"],
-        "dims": (node.kwargs["dim"],),
-        "starts": (node.kwargs["start"],),
-        "stops": (node.kwargs["start"] + node.kwargs["length"],),
-        "steps": (1,),
+        "dim": node.kwargs["dim"],
+        "start": node.kwargs["start"],
+        "stop": node.kwargs["start"] + node.kwargs["length"],
+        "step": 1,
     }
     with node.graph.inserting_before(node):
         new_node = node.graph.call_function(slice_tensor, kwargs=kwargs)
