[Relax][ONNX] Add CastLike support and dynamic-k Trilu to expand backend coverage by Aharrypotter · Pull Request #19898 · apache/tvm

Aharrypotter · 2026-06-28T13:29:45Z

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

# 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:

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:

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:

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: [Relax][ONNX] Use ONNX Backend Tests to improve frontend coverage.

- Add a CastLike converter to the Relax ONNX frontend (opset 15+). - CastLike casts the first input to the dtype of the second input, which is required by the expanded form of Relu in opset 18. - Add relu to the backend test allowlist. - Add a focused CastLike regression test. Relates to issue apache#19505.

- Remove the constant-k restriction from the Trilu converter. - For constant k, keep using the optimized relax.op.tril/triu. - For dynamic k, build the lower/upper-triangular mask explicitly with arange and broadcast it to the input shape. - Add tril and triu to the backend test allowlist. - Add a regression test for Trilu with dynamic k. Relates to issue apache#19505.

…ensions

gemini-code-assist

Code Review

This pull request introduces the CastLike ONNX operator converter and updates the Trilu converter to support dynamic k values by building a mask explicitly. It also expands the ONNX backend test suite to include several new operators. The review feedback suggests making the CastLike type retrieval safer to prevent potential AttributeErrors, removing redundant explicit broadcasts in the dynamic k mask generation, and using x.ty.dtype.dtype instead of x.ty.dtype for consistency when calling relax.const.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-06-28T13:33:31Z

+    def _impl_v15(cls, bb, inputs, attr, params):
+        data = inputs[0]
+        target = inputs[1]
+        target_dtype = target.ty.dtype.dtype


Accessing target.ty.dtype.dtype directly assumes that target.ty is populated and is a TensorType. If target is a relax.Var without type annotation or has a different type structure, this will raise an AttributeError. Consider using a safer fallback like getattr to handle potentially missing type information, similar to how it is done in the Cast converter.

gemini-code-assist · 2026-06-28T13:33:31Z

+        diff = relax.op.subtract(
+            relax.op.broadcast_to(col_idx, (m, n)),
+            relax.op.broadcast_to(row_idx, (m, n)),
+        )


The explicit relax.op.broadcast_to calls on col_idx and row_idx are redundant. In TVM Relax, element-wise binary operators like relax.op.subtract support implicit broadcasting natively. Removing these explicit broadcasts simplifies the generated Relax graph and avoids creating unnecessary intermediate operators.

diff = relax.op.subtract(col_idx, row_idx)

gemini-code-assist · 2026-06-28T13:33:31Z

-            return relax.op.tril(x, k)
+            mask = relax.op.less_equal(diff, k_int64)
+        mask = relax.op.broadcast_to(mask, shape)
+        return relax.op.where(mask, x, relax.const(0, x.ty.dtype))


To be consistent with the rest of the file (e.g., lines 372, 383), you should use x.ty.dtype.dtype to retrieve the string representation of the data type for relax.const, rather than passing the DataType object x.ty.dtype directly.

Suggested change

return relax.op.where(mask, x, relax.const(0, x.ty.dtype))

return relax.op.where(mask, x, relax.const(0, x.ty.dtype.dtype))

Aharrypotter added 3 commits June 28, 2026 20:04

[Relax][ONNX] Keep backend test registration open for model-level ext…

dab4f74

…ensions

gemini-code-assist Bot reviewed Jun 28, 2026

View reviewed changes

Aharrypotter added 2 commits June 28, 2026 21:35

[Relax][ONNX] Fix UP038 lint warnings in backend test adapter

2acccfb

[Relax][ONNX] Apply review feedback to CastLike and Trilu

a73ff50

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Relax][ONNX] Add CastLike support and dynamic-k Trilu to expand backend coverage#19898

[Relax][ONNX] Add CastLike support and dynamic-k Trilu to expand backend coverage#19898
Aharrypotter wants to merge 5 commits into
apache:mainfrom
Aharrypotter:onnx-castlike-relu

Aharrypotter commented Jun 28, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jun 28, 2026

Uh oh!

gemini-code-assist Bot Jun 28, 2026

Uh oh!

gemini-code-assist Bot Jun 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

	return relax.op.where(mask, x, relax.const(0, x.ty.dtype))
	return relax.op.where(mask, x, relax.const(0, x.ty.dtype.dtype))

Uh oh!

Conversation

Aharrypotter commented Jun 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Goal

What changed

Result

Design

CastLike support

Trilu dynamic k

Updated Allowlist

Safety Checks

Out of Scope / Non-Goals

Tests

References

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 28, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Aharrypotter commented Jun 28, 2026 •

edited

Loading

Trilu dynamic `k`