This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 7ff8d08b6c [Fix][Relax][Frontend][Torch] Support `aten.diagonal` from
decomposed repeated-subscript einsum (#20237)
7ff8d08b6c is described below
commit 7ff8d08b6c5f0433ec18d16caadee2590f4f302a
Author: HuEnwei <[email protected]>
AuthorDate: Fri Sep 11 03:10:34 2026 +0800
[Fix][Relax][Frontend][Torch] Support `aten.diagonal` from decomposed
repeated-subscript einsum (#20237)
Fixes: #20228
## Summary
`from_exported_program` runs `exported_program.run_decompositions()` by
default, and PyTorch's decomposition lowers `torch.einsum` with repeated
subscripts (diagonal / trace, e.g. `"ii->i"`, `"ii->"`, `"...ii->...i"`)
to
`aten.diagonal` + `permute` (+ `sum` for the trace).
`aten.diagonal.default`
was missing from the torch frontend `convert_map`, so **every** such
valid
model failed with:
```
AssertionError: Unsupported function types ['diagonal.default']
```
This PR adds an `aten.diagonal` converter and registers it in both the
exported-program and `from_fx` convert maps, so repeated-subscript
einsum —
and the directly-affected ops `torch.diagonal` / `torch.trace` — convert
and
run. Verified failing equations from the issue all convert with
`max|diff| = 0`
vs PyTorch.
## Root cause
`BaseFXGraphImporter._check_unsupported_func_type` asserts when a
`call_function` node's target is not in `convert_map`. For the einsum
family
above, `run_decompositions` introduces `aten.diagonal.default` nodes
that the
torch frontend had no handler for, so conversion aborts at the
assertion. This
is the same root cause for the direct ops `torch.diagonal` (lowered to
`diagonal.default` as-is) and `torch.trace` (lowered to `diagonal` +
`clone` +
`sum`). Skipping decomposition (`run_ep_decomposition=False`) keeps the
einsum
node intact and works — confirming the defect is the missing `diagonal`
handling, not `relax.op.einsum` semantics.
## Fix
Add `BaseFXGraphImporter._diagonal` in `base_fx_graph_translator.py`,
lowering
`diagonal(input, offset=0, dim1=0, dim2=1)` as:
1. `relax.op.permute_dims` — move `dim1` / `dim2` to the trailing two
axes;
2. two `relax.op.strided_slice` — crop each trailing axis to the
diagonal
length `min(extent1, extent2 ± offset)` (offset-adjusted), so the two
trailing extents are equal;
3. `relax.op.einsum([x], "...zz->...z")` — the repeated `z` label runs
over
both trailing axes simultaneously, extracting the diagonal.
The lowering handles static and dynamic (symbolic) shapes,
positive/negative
offsets, and arbitrary `dim1` / `dim2` (including negative indices).
Register
`"diagonal.default"` in `ExportedProgramImporter.create_convert_map` and
`"diagonal"` in `TorchFXImporter.create_convert_map`.
## Validation
### In-tree regression test (added)
`test_einsum_repeated_subscript` in
`tests/python/relax/test_frontend_from_exported_program.py`:
- `verify_model` against the exact lowering IR for `"ii->i"` on the
default
decomposition path (this case used to raise the assertion);
- `verify_model_numerically` for `"ii->"` (trace), `"...ii->...i"`
(batched
diagonal), the attention-style two-operand `"abca,abcb->c"`, and the
direct
ops `torch.diagonal(x, offset, 0, 1)` and `torch.trace`.
### Differential test
`verify_patch.py` runs on the locked build and simulates the pre-fix
behavior
at runtime (popping `diagonal.default` from the generated convert map):
- **Baseline (pre-fix)**: all 21 diagonal-producing cases (10 issue
einsum
equations + 11 direct `torch.diagonal`/`torch.trace`/`torch.diag`)
reproduce
the exact `AssertionError: Unsupported function types
['diagonal.default']`;
1 case (`torch.diag` on a 1-D input, which goes through `diag_embed`) is
unaffected and stays correct in baseline.
- **Post-fix**: all 22 issue + direct-op cases convert and match PyTorch
with
`max|diff| = 0`.
- **Dynamic shapes**: `"ii->i"` and `"...ii->...i"` with symbolic dims
(both
diagonal dims sharing one `Dim`) match PyTorch exactly.
- **Regression**: the regular einsum family (matmul, transpose, dot,
outer,
batch matmul, ellipsis broadcasting/summation, 3-operand, implicit
output) —
15 cases — all still match with `max|diff| = 0`.
Run:
```bash
TVM_LIBRARY_PATH=<tvm>/build/lib PYTHONPATH=<tvm 源码>/python \
/home/shenqingchao/miniconda3/envs/tvm23/bin/python \
results/TVM/deepseek-v4-flash/prove_hum/torch_einsum/verify_patch.py
```
## Files changed
- `python/tvm/relax/frontend/torch/base_fx_graph_translator.py` — add
`_diagonal` (permute_dims + strided_slice crop + einsum `...zz->...z`).
- `python/tvm/relax/frontend/torch/exported_program_translator.py` —
register
`"diagonal.default"` in the exported-program `convert_map`.
- `python/tvm/relax/frontend/torch/fx_translator.py` — register
`"diagonal"` in
the `from_fx` `convert_map`.
- `tests/python/relax/test_frontend_from_exported_program.py` — add
`test_einsum_repeated_subscript` regression coverage.
---
.../frontend/torch/base_fx_graph_translator.py | 101 ++++++++++++++
.../frontend/torch/exported_program_translator.py | 1 +
python/tvm/relax/frontend/torch/fx_translator.py | 1 +
.../relax/test_frontend_from_exported_program.py | 151 +++++++++++++++++++++
4 files changed, 254 insertions(+)
diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
index f6781f026f..d42590c256 100644
--- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
+++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py
@@ -34,6 +34,29 @@ from tvm.ir import PrimType
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 @@ class BaseFXGraphImporter(metaclass=abc.ABCMeta):
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,
diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py
b/python/tvm/relax/frontend/torch/exported_program_translator.py
index ced0aa7b28..e6df019c5c 100644
--- a/python/tvm/relax/frontend/torch/exported_program_translator.py
+++ b/python/tvm/relax/frontend/torch/exported_program_translator.py
@@ -1894,6 +1894,7 @@ class ExportedProgramImporter(BaseFXGraphImporter):
"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]]
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py
b/python/tvm/relax/frontend/torch/fx_translator.py
index bef0b58f08..4517650112 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -952,6 +952,7 @@ class TorchFXImporter(BaseFXGraphImporter):
"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,
diff --git a/tests/python/relax/test_frontend_from_exported_program.py
b/tests/python/relax/test_frontend_from_exported_program.py
index 53e8a3a314..75962826f9 100644
--- a/tests/python/relax/test_frontend_from_exported_program.py
+++ b/tests/python/relax/test_frontend_from_exported_program.py
@@ -3327,6 +3327,157 @@ def test_einsum():
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):
