feat: support aten.diagonal converter (#2856)

Author: chohk88

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -3036,6 +3036,64 @@ def aten_ops_flip(
     )
 
 
+def zero_diag_size_validator(node: Node) -> bool:
+    meta = node.args[0].meta.get("tensor_meta")
+    if meta:
+        input_shape = meta.shape
+    else:
+        _LOGGER.warning(
+            "Meta information of input is missing. Unable to validate diagonal size, falling back to PyTorch operation."
+        )
+        return False
+
+    offset, dim1, dim2 = (
+        node.args[1],
+        node.args[2],
+        node.args[3],
+    )
+
+    num_dims = len(input_shape)
+
+    # Adjust dimensions to be positive and canonicalize
+    dim1 = get_positive_dim(dim1, num_dims)
+    dim2 = get_positive_dim(dim2, num_dims)
+
+    if offset >= 0:
+        diag_size = max(min(input_shape[dim1], input_shape[dim2] - offset), 0)
+    else:
+        diag_size = max(min(input_shape[dim1] + offset, input_shape[dim2]), 0)
+
+    if diag_size == 0:
+        _LOGGER.debug(
+            "Diagonal size is zero, resulting in an empty tensor which is not supported for this operation."
+        )
+        return False
+    else:
+        return True
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten.diagonal.default, capability_validator=zero_diag_size_validator
+)
+def aten_ops_diagonal(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.slice.diagonal(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args_bounds_check(args, 1, replacement=0),
+        args_bounds_check(args, 2, replacement=0),
+        args_bounds_check(args, 3, replacement=1),
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.scalar_tensor.default)
 def aten_ops_scalar_tensor(
     ctx: ConversionContext,

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -653,3 +653,34 @@ def set_item(
         0,
     )
     return ans
+
+
+def calculate_strides(shape: Sequence[int]) -> Sequence[int]:
+    """
+    Calculate the strides for a given shape of a multi-dimensional array.
+
+    The output stride for each dimension indicates the number of elements to skip in
+    memory to move to the next element along that dimension. The last dimension always
+    has a stride of 1 because elements are stored contiguously along this dimension.
+
+    Example:
+        For a 3-dimensional array with shape [2, 3, 4]:
+        - shape = [2, 3, 4]
+        - The function will calculate the strides as follows:
+            1. Initialize strides: [1, 1, 1]
+            2. Calculate strides for each dimension from right to left:
+               - For i = 1: strides[1] = strides[2] * shape[2] = 1 * 4 = 4
+               - For i = 0: strides[0] = strides[1] * shape[1] = 4 * 3 = 12
+            - Final strides: [12, 4, 1]
+
+        Therefore, the output will be [12, 4, 1].
+
+        This means:
+        - To move along the first dimension, skip 12 elements.
+        - To move along the second dimension, skip 4 elements.
+        - To move along the third dimension, skip 1 element.
+    """
+    strides = [1] * len(shape)
+    for i in range(len(shape) - 2, -1, -1):
+        strides[i] = strides[i + 1] * shape[i + 1]
+    return strides

py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py

@@ -8,6 +8,7 @@
 from torch_tensorrt.dynamo.conversion import impl
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
 from torch_tensorrt.dynamo.conversion.converter_utils import (
+    calculate_strides,
     flatten_dims,
     get_positive_dim,
     get_trt_tensor,
@@ -262,6 +263,68 @@ def flip(
     return layer.get_output(0)
 
 
+def diagonal(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    offset: int,
+    dim1: int,
+    dim2: int,
+) -> TRTTensor:
+    """
+    This implementation is inspired by the reference implementation in PyTorch:
+    https://github.com/pytorch/pytorch/blob/082251e76b93b277ff2791d0e2b64934add34644/torch/_refs/__init__.py#L4255
+    """
+    input_shape = input.shape
+    num_dims = len(input_shape)
+
+    # Adjust dimensions to be positive and canonicalize
+    dim1 = get_positive_dim(dim1, num_dims)
+    dim2 = get_positive_dim(dim2, num_dims)
+
+    # Calculate the size of the diagonal
+    if offset >= 0:
+        diag_size = max(min(input_shape[dim1], input_shape[dim2] - offset), 0)
+    else:
+        diag_size = max(min(input_shape[dim1] + offset, input_shape[dim2]), 0)
+
+    if diag_size == 0:
+        raise ValueError("The size of the diagonal is non-positive.")
+
+    strides = calculate_strides(input_shape)
+
+    # Compute the storage offset
+    storage_offset = 0
+    if offset >= 0:
+        storage_offset += offset * strides[dim2]
+    else:
+        storage_offset -= offset * strides[dim1]
+
+    # Calculate new sizes and strides for as_strided
+    sizes = [s for i, s in enumerate(input_shape) if i not in (dim1, dim2)]
+    sizes.append(diag_size)
+
+    input_strides = [s for i, s in enumerate(strides) if i not in (dim1, dim2)]
+    new_stride = strides[dim1] + strides[dim2]
+    input_strides.append(new_stride)
+
+    # Use as_strided to get the diagonal elements
+    diagonal_output = as_strided(
+        ctx,
+        target,
+        source_ir,
+        f"{name}_as_strided",
+        input,
+        sizes,
+        input_strides,
+        storage_offset,
+    )
+
+    return diagonal_output
+
+
 def as_strided(
     ctx: ConversionContext,
     target: Target,

tests/py/dynamo/conversion/test_diagonal_aten.py

@@ -0,0 +1,124 @@
+import torch
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
+
+from .harness import DispatchTestCase
+
+
+class TestAsStridedConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            (
+                (3, 3),
+                1,
+                0,
+                1,
+            ),
+            (
+                (3, 3),
+                1,
+                0,
+                -1,
+            ),
+            (
+                (3, 4),
+                1,
+                0,
+                1,
+            ),
+            (
+                (5, 4, 2),
+                -1,
+                1,
+                2,
+            ),
+            (
+                (5, 4, 2),
+                1,
+                2,
+                0,
+            ),
+            (
+                (6, 5, 4),
+                1,
+                0,
+                1,
+            ),
+            (
+                (2, 5, 4, 2),
+                0,
+                0,
+                1,
+            ),
+            (
+                (2, 5, 4, 2),
+                1,
+                1,
+                2,
+            ),
+            (
+                (2, 5, 4, 2),
+                1,
+                -1,
+                2,
+            ),
+            (
+                (2, 5, 4, 2),
+                1,
+                1,
+                -2,
+            ),
+            (
+                (2, 5, 4, 2),
+                1,
+                -1,
+                -2,
+            ),
+            (
+                (2, 5, 4, 2),
+                0,
+                0,
+                2,
+            ),
+            (
+                (2, 5, 4, 2),
+                -1,
+                1,
+                2,
+            ),
+            (
+                (2, 5, 4, 2, 6),
+                1,
+                1,
+                2,
+            ),
+            (
+                (2, 5, 4, 2, 5, 6),
+                1,
+                1,
+                2,
+            ),
+        ]
+    )
+    def test_diagonal(
+        self,
+        input_shape,
+        offset,
+        dim1,
+        dim2,
+    ):
+        class TestModule(torch.nn.Module):
+            def forward(self, x):
+                return torch.ops.aten.diagonal.default(x, offset, dim1, dim2)
+
+        inputs = [torch.randn(input_shape)]
+        self.run_test(
+            TestModule(),
+            inputs,
+            enable_passes=True,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()