[Fix][Relax][Frontend][Torch] Support aten.diagonal from decomposed repeated-subscript einsum

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -34,6 +34,29 @@
 from tvm.runtime import DataTypeCode
 
 
+def _diagonal_einsum_subscripts(ndim: int, dim1: int, dim2: int) -> str:
+    """Return explicit einsum subscripts that extract the ``dim1``/``dim2`` diagonal.
+
+    This is the fast-path lowering for :meth:`BaseFXGraphImporter._diagonal`
+    when ``offset == 0`` and both diagonal axes have the same extent. The
+    non-diagonal axes keep their natural order in the output while the diagonal
+    axis is appended last, matching ``aten.diagonal``. Non-diagonal axes are
+    labelled ``a``, ``b``, ... and the repeated (diagonal) label is ``z``, e.g.
+    an ``N x N`` input with ``dim1 == 0``, ``dim2 == 1`` gives ``"zz->z"``.
+    """
+    labels = [None] * ndim
+    # Non-diagonal axes are labelled from ``a`` to ``y``; ``z`` is reserved for
+    # the repeated (diagonal) label so the two never collide.
+    letters = iter(ch for ch in "abcdefghijklmnopqrstuvwxyz" if ch != "z")
+    for i in range(ndim):
+        if i != dim1 and i != dim2:
+            labels[i] = next(letters)
+    labels[dim1] = "z"
+    labels[dim2] = "z"
+    leading = "".join(labels[i] for i in range(ndim) if i != dim1 and i != dim2)
+    return f"{''.join(labels)}->{leading}z"
+
+
 class BaseFXGraphImporter(metaclass=abc.ABCMeta):
     """Base class for FX Graph Importer."""
 
@@ -1264,6 +1287,84 @@ def _einsum(self, node: fx.Node) -> relax.Var:
         operands = args[1] if isinstance(args[1], torch.Size | tuple | list) else args[1:]
         return self.block_builder.emit(relax.op.einsum(operands, args[0]))
 
+    def _diagonal(self, node: fx.Node) -> relax.Var:
+        """Convert ``aten.diagonal`` / ``torch.diagonal`` to Relax.
+
+        ``diagonal(input, offset=0, dim1=0, dim2=1)`` extracts the elements
+        ``input[..., i, i + offset]`` along the ``dim1`` / ``dim2`` axes. It
+        shows up in the exported graph through ``run_decompositions`` of
+        ``torch.einsum`` with repeated subscripts (e.g. ``"ii->i"``,
+        ``"ii->"``, ``"...ii->...i"``), which lower to an ``aten.diagonal``
+        followed by a ``sum`` reduction.
+
+        We lower it as: when ``offset == 0`` and both diagonal axes have equal
+        extent (the common ``torch.einsum("ii->i")`` case), a single einsum
+        whose subscript repeats one label over the two axes extracts the
+        diagonal directly, so no full-size permute / slice intermediate is
+        materialized. Otherwise we permute ``dim1`` / ``dim2`` to the trailing
+        two axes, slice each trailing axis to the diagonal length (min of the
+        two extents, adjusted by ``offset``), and take the diagonal with an
+        einsum contraction ``...zz->...z`` (the repeated ``z`` label runs over
+        both trailing axes simultaneously).
+        """
+
+        args = self.retrieve_args(node)
+        x = args[0]
+        offset = args[1] if len(args) > 1 else node.kwargs.get("offset", 0)
+        dim1 = args[2] if len(args) > 2 else node.kwargs.get("dim1", 0)
+        dim2 = args[3] if len(args) > 3 else node.kwargs.get("dim2", 1)
+
+        shape = self.shape_of(x)
+        ndim = len(shape.values)
+        dim1 = dim1 if dim1 >= 0 else ndim + dim1
+        dim2 = dim2 if dim2 >= 0 else ndim + dim2
+        if dim1 == dim2:
+            raise ValueError(f"diagonal requires dim1 != dim2, got {dim1} == {dim2}")
+
+        offset = int(offset)
+
+        n = shape.values[dim1]
+        m = shape.values[dim2]
+        # Fast path for the common ``offset == 0`` case with equal extents on the
+        # diagonal axes (e.g. torch.einsum("ii->i") on an N x N input). The
+        # diagonal can be read with a single einsum that repeats one subscript
+        # label over the two axes (``relax.op.einsum`` lowers it to one O(N)
+        # loop), avoiding the identity permute / strided-slice that would each
+        # materialize a full-size O(N^2) copy. Non-diagonal axes must still fit
+        # in the single-letter einsum label alphabet.
+        if offset == 0 and ndim - 2 <= 25 and tvm_ffi.structural_equal(n, m):
+            subscripts = _diagonal_einsum_subscripts(ndim, dim1, dim2)
+            return self.block_builder.emit(relax.op.einsum([x], subscripts))
+
+        # Move dim1, dim2 to the trailing two axes.
+        perm = [i for i in range(ndim) if i != dim1 and i != dim2] + [dim1, dim2]
+        permuted = self.block_builder.emit(relax.op.permute_dims(x, perm))
+
+        if offset >= 0:
+            diag_len = tirx.max(0, tirx.min(n, m - offset))
+            begin1, end1 = 0, diag_len
+            begin2, end2 = offset, offset + diag_len
+        else:
+            diag_len = tirx.max(0, tirx.min(n + offset, m))
+            begin1, end1 = -offset, -offset + diag_len
+            begin2, end2 = 0, diag_len
+
+        # Crop both diagonal axes to the diagonal length so the einsum ``z``
+        # label sees equal extents on both trailing axes.
+        cropped = self.block_builder.emit(
+            relax.op.strided_slice(
+                permuted, axes=[ndim - 2], begin=[begin1], end=[end1], strides=[1]
+            )
+        )
+        cropped = self.block_builder.emit(
+            relax.op.strided_slice(
+                cropped, axes=[ndim - 1], begin=[begin2], end=[end2], strides=[1]
+            )
+        )
+
+        # ``...zz -> ...z``: keep every leading axis, contract the diagonal pair.
+        return self.block_builder.emit(relax.op.einsum([cropped], "...zz->...z"))
+
     def _embedding_impl(
         self,
         x,

python/tvm/relax/frontend/torch/exported_program_translator.py

@@ -1894,6 +1894,7 @@ def create_convert_map(
             "conv3d.default": self._conv3d,
             "convolution.default": self._convolution,
             "cross_entropy_loss.default": self._cross_entropy_default,
+            "diagonal.default": self._diagonal,
             "einsum.default": self._einsum,
             "embedding.default": lambda node: self._embedding_impl(
                 self.env[node.args[1]], self.env[node.args[0]]

python/tvm/relax/frontend/torch/fx_translator.py

@@ -952,6 +952,7 @@ def create_convert_map(
             "conv2d": self._conv2d,
             "conv3d": self._conv3d,
             "cross_entropy": self._cross_entropy,
+            "diagonal": self._diagonal,
             "einsum": self._einsum,
             "interpolate": self._interpolate,
             "layer_norm": self._layer_norm,

tests/python/relax/test_frontend_from_exported_program.py

@@ -3327,6 +3327,157 @@ def main(
     verify_model(Einsum2(), example_args, {}, Expected2, run_ep_decomposition=False)
 
 
+def test_einsum_repeated_subscript():
+    """einsum with repeated subscripts (diagonal / trace) on the default
+    decomposition path.
+
+    ``run_decompositions`` (default) lowers repeated-subscript einsum to
+    ``aten.diagonal`` + ``permute`` (+ ``sum`` for the trace), which the
+    frontend converts with the ``_diagonal`` lowering. For the zero-offset
+    square case (e.g. ``torch.einsum("ii->i")`` on an ``N x N`` input) the
+    frontend emits a single repeated-subscript einsum that reads the diagonal
+    directly; otherwise it permutes the diagonal dims to the trailing two axes,
+    slices each to the diagonal length, and runs an einsum ``...zz->...z``.
+    This used to raise ``AssertionError: Unsupported function types
+    ['diagonal.default']``.
+    """
+
+    class EinsumDiag(Module):
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, x):
+            return torch.einsum("ii->i", x)
+
+    @tvm.script.ir_module
+    class Expected:
+        @R.function
+        def main(x: R.Tensor((3, 3), dtype="float32")) -> R.Tuple(R.Tensor((3,), dtype="float32")):
+            with R.dataflow():
+                lv: R.Tensor((3,), dtype="float32") = R.einsum((x,), subscripts="zz->z")
+                lv1: R.Tensor((3,), dtype="float32") = R.permute_dims(lv, axes=[0])
+                lv2: R.Tensor((3,), dtype="float32") = R.permute_dims(lv1, axes=[0])
+                gv: R.Tuple(R.Tensor((3,), dtype="float32")) = (lv2,)
+                R.output(gv)
+            return gv
+
+    example_args = (torch.randn(3, 3, dtype=torch.float32),)
+    verify_model(EinsumDiag(), example_args, {}, Expected)
+
+    class TraceEinsum(Module):
+        def forward(self, x):
+            return torch.einsum("ii->", x)
+
+    class BatchedDiagEinsum(Module):
+        def forward(self, x):
+            return torch.einsum("...ii->...i", x)
+
+    class AttentionEinsum(Module):
+        def forward(self, x, y):
+            return torch.einsum("abca,abcb->c", x, y)
+
+    verify_model_numerically(TraceEinsum(), (torch.randn(4, 4),))
+    verify_model_numerically(BatchedDiagEinsum(), (torch.randn(2, 3, 3),))
+    verify_model_numerically(AttentionEinsum(), (torch.randn(3, 3, 4, 3), torch.randn(3, 3, 4, 3)))
+
+    class DirectDiagonal(Module):
+        def __init__(self):
+            super().__init__()
+            self.offset = 1
+
+        def forward(self, x):
+            return torch.diagonal(x, self.offset, 0, 1)
+
+    class DirectTrace(Module):
+        def forward(self, x):
+            return torch.trace(x)
+
+    verify_model_numerically(DirectDiagonal(), (torch.randn(3, 4),))
+    verify_model_numerically(DirectTrace(), (torch.randn(4, 4),))
+
+    # Out-of-range offsets (|offset| >= max(extent1, extent2)) are valid in
+    # PyTorch and yield an empty diagonal of shape (0,); the lowering must
+    # clamp the diagonal length to zero instead of producing negative slice
+    # extents or a wrong non-empty shape.
+    class DirectDiagonalOutOfRange(Module):
+        def __init__(self, offset):
+            super().__init__()
+            self.offset = offset
+
+        def forward(self, x):
+            return torch.diagonal(x, self.offset, 0, 1)
+
+    for offset in [4, 5, 6, -3, -4, -5, -6]:
+        verify_model_numerically(DirectDiagonalOutOfRange(offset), (torch.randn(3, 4),))
+
+
+def test_einsum_diagonal_lowers_without_full_size_intermediate():
+    """Regression test: a zero-offset square diagonal must not materialize
+    full-size intermediates.
+
+    ``torch.einsum("ii->i")`` on an ``N x N`` input is decomposed to
+    ``aten.diagonal`` by ``run_decompositions``. Lowering that diagonal by
+    permuting the diagonal dims to the trailing axes, slicing each to the
+    diagonal length, and running the ``...zz->...z`` einsum materializes three
+    full-size ``N x N`` intermediates (an identity permute and two identity
+    strided slices) and hence three O(N^2) copy loops before the final O(N)
+    diagonal loop. The ``_diagonal`` fast path instead emits a single
+    repeated-subscript einsum that reads the diagonal directly, so no full-size
+    intermediate exists in the frontend graph (and therefore neither in the
+    lowered TIR). Assert that every intermediate produced by a call is at most
+    O(N), both before and after legalization.
+    """
+
+    class EinsumDiag(Module):
+        def forward(self, x):
+            return torch.einsum("ii->i", x)
+
+    n = 8
+    exported_program = export(EinsumDiag(), args=(torch.randn(n, n),))
+    mod = from_exported_program(exported_program)
+
+    def rank2_call_results(ir_mod):
+        """Names of calls whose result is a rank-2 (full-size) tensor."""
+        results = []
+        for func in ir_mod.functions.values():
+            if not isinstance(func, relax.Function):
+                continue
+            for block in func.body.blocks:
+                for binding in block.bindings:
+                    if not (
+                        isinstance(binding.value, relax.Call)
+                        and isinstance(binding.value.op, tvm.ir.Op)
+                    ):
+                        continue
+                    if isinstance(binding.var.ty, relax.TensorType) and binding.var.ty.ndim == 2:
+                        results.append(binding.value.op.name)
+        return results
+
+    # The diagonal must be the only full-size (N x N) tensor touched: it is the
+    # function input read directly by a single repeated-subscript einsum. No
+    # call may produce a rank-2 intermediate.
+    assert rank2_call_results(mod) == []
+
+    # Sanity check that the graph really performs the diagonal: exactly one
+    # einsum on the N x N input producing an N-vector.
+    einsum_calls = []
+    for block in mod["main"].body.blocks:
+        for binding in block.bindings:
+            if (
+                isinstance(binding.value, relax.Call)
+                and isinstance(binding.value.op, tvm.ir.Op)
+                and binding.value.op.name == "relax.einsum"
+            ):
+                einsum_calls.append(binding.var)
+    assert len(einsum_calls) == 1
+    assert einsum_calls[0].ty.ndim == 1
+
+    # Legalize and check again on the lowered graph.
+    with tvm.target.Target("llvm"):
+        lowered = relax.transform.LegalizeOps()(mod)
+    assert rank2_call_results(lowered) == []
+
+
 def test_outer():
     class Outer(torch.nn.Module):
         def forward(self, x, y):

tests/python/relax/test_frontend_onnx.py

@@ -1230,7 +1230,7 @@ def test_multi_input_unknown_static_shape(op_name, num_inputs):
         "Slice", ["x", "starts", "ends", "axes"], ["sliced"], name="slice0"
     )
     other_names = [f"y{i}" for i in range(num_inputs - 1)]
-    op_node = helper.make_node(op_name, ["sliced"] + other_names, ["output"], name="op0")
+    op_node = helper.make_node(op_name, ["sliced", *other_names], ["output"], name="op0")
 
     graph = helper.make_graph(
         [slice_node, op_node],
@@ -1241,10 +1241,7 @@ def test_multi_input_unknown_static_shape(op_name, num_inputs):
             helper.make_tensor_value_info("ends", TensorProto.INT64, [1]),
             helper.make_tensor_value_info("axes", TensorProto.INT64, [1]),
         ]
-        + [
-            helper.make_tensor_value_info(name, TensorProto.FLOAT, [2, 3])
-            for name in other_names
-        ],
+        + [helper.make_tensor_value_info(name, TensorProto.FLOAT, [2, 3]) for name in other_names],
         outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT, [2, 3])],
     )
     model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])