Aharrypotter opened a new pull request, #19898:
URL: https://github.com/apache/tvm/pull/19898
## Summary
This PR adds `CastLike` support and dynamic-`k` support for `Trilu` in the
Relax
ONNX frontend, then adds `relu`, `tril`, and `triu` to the official ONNX
backend
test allowlist.
### Goal
Increase the Relax ONNX frontend's coverage in the official ONNX Backend Test
Suite by enabling operators that already have hand-written frontend tests but
still fail some official node-level tests.
### What changed
- Added a `CastLike` converter.
- Removed the constant-`k` restriction from the `Trilu` converter.
- Added `relu`, `tril`, and `triu` to `_INCLUDE_OPS`.
- Added `_EXCLUDE_PATTERNS` to filter out a few model-level tests whose names
collide with the node-level include patterns.
### Result
```text
# Before
388 passed, 3142 skipped
# After
451 passed, 3377 skipped
```
This PR directly addresses part of #19505.
## Design
### CastLike support
ONNX `CastLike` (opset 15+) takes two inputs: the data to cast and a tensor
whose dtype determines the output dtype. The opset-18 expanded form of `Relu`
decomposes the operator into a subgraph that uses `CastLike`, so importing
any
opset-18 `Relu` model previously failed with:
```text
OpNotImplemented: The following operators are not supported for frontend
ONNX: CastLike
```
The new `CastLike` converter reads the dtype of the second input and emits
`relax.op.astype(data, target_dtype)`. It handles both constant and dynamic
target tensors because the dtype is taken from the input's type information.
### Trilu dynamic `k`
The existing `Trilu` converter only accepted a constant `k` diagonal offset
and
raised `ValueError` for any dynamic / graph-input `k`. Several official ONNX
node tests (`test_tril_neg`, `test_triu_zero`, etc.) supply `k` as a graph
input, so those tests could not pass.
The converter now branches:
- If `k` is a constant or omitted, use the optimized `relax.op.tril` /
`relax.op.triu` paths.
- If `k` is dynamic, construct the lower/upper-triangular mask explicitly:
1. Build row and column index tensors with `relax.op.arange`.
2. Compute `col_index - row_index`.
3. Compare against the dynamic scalar `k`.
4. Broadcast the mask to the input shape and use `relax.op.where` to zero
the excluded elements.
## Updated Allowlist
| Operator | Added to `_INCLUDE_OPS` | Tests gained |
|---|---|---|
| `relu` | yes | 2 |
| `tril` | yes | 18 |
| `triu` | yes | 18 |
Total backend suite progress: **388 passed → 451 passed** (all CPU; CUDA
tests
are registered but skipped because the backend adapter only supports CPU).
## Safety Checks
- `CastLike` returns `relax.op.astype(data, target_dtype)` where
`target_dtype` is the dtype of the second input.
- Constant / omitted `k` in `Trilu` keeps the existing optimized
`relax.op.tril` / `relax.op.triu` lowering.
- Dynamic `k` in `Trilu` is implemented without calling `relax.op.tril` /
`triu` with a non-constant diagonal offset.
- `_INCLUDE_OPS` remains the gate for which backend tests run; a small
`_EXCLUDE_PATTERNS` list filters model-level name collisions so the suite
stays green without limiting the registered test classes.
## Out of Scope / Non-Goals
- This PR does not address the other candidate operators that still fail node
tests (`cast`, `equal`, `gather`, `reshape`, `shape`, `reduce_*`). Those
will
be handled in follow-up PRs.
- This PR does not change the frontend's handling of `Relu` itself; it only
unblocks the expanded form by adding `CastLike`.
- This PR does not add CUDA support to the backend test adapter.
## Tests
| Test | Coverage |
|---|---|
| `test_castlike_ir` | New `CastLike` converter, structural IR check |
| `test_trilu` / `test_trilu_with_const_k` | Existing Trilu coverage,
unchanged |
| `test_trilu_dynamic_k_ir` | New parametrized structural IR test for
dynamic `k` (`upper=True/False`) |
| `test_frontend_onnx_backend.py` | Official ONNX node tests for `relu`,
`tril`, `triu` |
Local validation:
```bash
python -m pytest tests/python/relax/test_frontend_onnx.py::test_castlike_ir
-xvs
python -m pytest tests/python/relax/test_frontend_onnx.py -k "trilu" -xvs
python -m pytest tests/python/relax/test_frontend_onnx_backend.py -q
python -m ruff format --check \
python/tvm/relax/frontend/onnx/onnx_frontend.py \
tests/python/relax/test_frontend_onnx.py \
tests/python/relax/test_frontend_onnx_backend.py
python -m ruff check \
python/tvm/relax/frontend/onnx/onnx_frontend.py \
tests/python/relax/test_frontend_onnx.py \
tests/python/relax/test_frontend_onnx_backend.py
```
Result:
```text
test_castlike_ir: passed
test_frontend_onnx.py -k "trilu": 10 passed
test_frontend_onnx_backend.py -q: 450 passed, 3080 skipped
ruff format --check: 3 files already formatted
ruff check: All checks passed
```
## References
- Relates to [#19505](https://github.com/apache/tvm/issues/19505):
`[Relax][ONNX]
Use ONNX Backend Tests to improve frontend coverage`.
- Follows the backend test positioning notes in
`.memory/info/relax-onnx-backend-tests-positioning.md`.
--
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]
