apache / apache/tvm

[RFC] Expanding ONNX operator coverage in the Relax frontend

Open
#20,124 1 comment 0 reactions 0 assignees View on GitHub
needs-triage type:rfc-tracking
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 1h
Merged PRs (30d)
112

Description

## Context

The Relax ONNX frontend currently implements ~160 operators from the ONNX spec. Comparing against the current operator list (https://onnx.ai/onnx/operators/index.html), a number of ops are not yet supported.

Before opening PRs, I'd like to check on priorities/approach with maintainers. I went through `python/tvm/relax/op/` and `python/tvm/topi/` to check what's actually available to build on, rather than guessing from op names alone, so the grouping below reflects what I found there.

## Missing operators

| Category | Operators |
|---|---|
| Activations / math | Celu, Swish, LpNormalization, LinearAttention, Det, Col2Im, TensorScatter |
| Casting / bit-level | CastLike, BitCast |
| Cropping / sequence | CenterCropPad, ReverseSequence, SequenceMap |
| Normalization | GroupNormalization |
| Recurrent networks | RNN, GRU, LSTM |
| Convolution | ConvInteger, DeformConv, CausalConvWithState |
| Quantization | QLinearConv, QLinearMatMul |
| Random / sampling | RandomNormal, RandomNormalLike, RandomUniform, RandomUniformLike, Multinomial, Bernoulli |
| Signal processing | DFT, STFT, BlackmanWindow, HammingWindow, HannWindow, MelWeightMatrix |
| Control flow | Loop, Scan |
| Loss functions | NegativeLogLikelihoodLoss, SoftmaxCrossEntropyLoss |
| String / text | StringConcat, StringSplit, StringNormalizer, RegexFullMatch, TfIdfVectorizer |
| Attention / misc | RotaryEmbedding, ImageDecoder |

## Implementation notes

These vary a lot in difficulty. Some already have a clear path based on what's in the codebase:

- **RNN, GRU, LSTM** - no fused recurrent op in Relax, but the torch and tflite Relax frontends already implement LSTM by unrolling the recurrence with existing ops; the ONNX converter could follow the same pattern.
- **CastLike, Swish** - Relax has direct equivalents (`astype`, `nn.silu`) that cover the common case (default `alpha=1.0` for Swish; standard dtype casts for CastLike). Non-default variants (float8 `round_mode`/`saturate` for CastLike, non-unit alpha for Swish) would need a small amount of extra composition.
- **CumProd, ReverseSequence** - Relax has ops covering the same purpose (`cumprod`, `reverse_sequence`), but attribute defaults/semantics differ from the ONNX spec (Relax's `cumprod` has no reverse mode; `reverse_sequence`'s axis defaults are swapped relative to ONNX's `time_axis`/`batch_axis`), so the converter needs explicit attribute mapping rather than a pass-through.

Others (quantized conv/matmul, signal processing, `Loop`/`Scan`) likely need new compute/schedule work or, in the case of `Loop`/`Scan`, graph-level control-flow support. I'm not certain of the right approach here and would appreciate maintainer input.

## Questions for maintainers

1. Is there an existing priority order for ONNX op coverage, or known user demand for any of the above from specific model families?
2. Are any of these considered out of scope or low priority for the Relax frontend (e.g. control flow, quantized ops, signal processing)?

## Proposal

I'd like to pick up work on some of these, starting with the ones that have a clear implementation path (CastLike, Swish, ReverseSequence, CumProd attribute mapping, then RNN/GRU/LSTM via unrolling).

## Questions for the maintainers

- Is there an existing priority order for ONNX op coverage, or known user demand for any of the above from specific model families?
- Are any of these (e.g. control flow, quantized ops, signal processing) considered out of scope or low priority for the Relax frontend?

Happy to pick up items above, split into smaller PRs by category, and help review PRs from others in this area as it progresses.

### Triage
* needs-triage
* frontend:onnx
* status: RFC

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reading the operator implementations and available building blocks in python/tvm/relax/op/ and python/tvm/topi/, then compare the proposed operators with the torch and tflite Relax frontend LSTM implementations. Focus first on the operators maintainers prioritize, such as CastLike, Swish, ReverseSequence, CumProd, or recurrent operators. Done means maintainers have agreed on scope and priority and the work has been split into focused, independently implementable PRs.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.