tlopex commented on PR #20237:
URL: https://github.com/apache/tvm/pull/20237#issuecomment-5530609421
Thanks for fixing the out-of-range offset case. One remaining concern is the
asymptotic cost of this lowering.
For an `N x N` input with `offset=0` (the common `torch.einsum("ii->i")`
case), this emits an identity `permute_dims`, two identity `strided_slice`
operations, and then the diagonal einsum. After `zero_pipeline`, the generated
TIR still contains three `N x N` intermediate buffers and three `O(N²)` copy
loops before the final `O(N)` diagonal loop. In comparison,
`relax.op.einsum([x], "ii->i")` lowers directly to a single `O(N)` loop.
This can make diagonal extraction unnecessarily slow and may cause OOM for
large inputs. Could we add a direct diagonal lowering, or at least a fast path
for `offset == 0` with equal diagonal dimensions, and add a test that checks
the lowered TIR does not materialize full-size intermediates?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
