[Relax][ONNX] Add Optional and MatMulInteger16 frontend support by Aharrypotter · Pull Request #18950

[Relax][ONNX] Add Optional and MatMulInteger16 frontend support by Aharrypotter · Pull Request #18950 · apache/tvm

Summary

This PR adds Relax ONNX frontend support for:

Optional
OptionalHasElement
OptionalGetElement
MatMulInteger16 from the com.microsoft domain

The implementation follows existing TVM ONNX frontend patterns and keeps Optional handling explicit through an empty-Optional sentinel during import.

Changes

add ONNX frontend converters for Optional, OptionalHasElement, and OptionalGetElement
add ONNX frontend converter for MatMulInteger16
extend ONNX attribute parsing to handle TYPE_PROTO
preserve empty Optional values during import and unwrap them consistently
register Optional-related ops and MatMulInteger16 in the ONNX converter map
handle Optional outputs correctly in importer output counting and normalization
tighten converter docstrings and input validation for better consistency with nearby TVM code

Tests

Added or updated tests in tests/python/relax/test_frontend_onnx.py to cover:

numerical correctness for MatMulInteger16
structural IR checks for MatMulInteger16
invalid dtype rejection for MatMulInteger16
tensor and sequence Optional round-trips
empty Optional behavior for OptionalHasElement
structural IR checks ensuring Optional ops are erased as expected
missing type attribute rejection for empty Optional
empty OptionalGetElement rejection

Validation

Validated with:

python -m ruff check python/tvm/relax/frontend/onnx/onnx_frontend.py tests/python/relax/test_frontend_onnx.py
python -m pytest -n 1 tests/python/relax/test_frontend_onnx.py -k "optional or matmulinteger16" -v

Result:

13 passed

This PR completes the ONNX MatMulInteger16 and Optional work tracked in #18945.

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds support for the MatMulInteger16, Optional, OptionalHasElement, and OptionalGetElement ONNX operators to the Relax frontend. It introduces a sentinel object for empty optionals, updates attribute parsing to handle type protocols, and includes comprehensive unit tests for the new functionality. Feedback suggests tightening input validation for OptionalHasElement to strictly adhere to the ONNX specification and improving test robustness by programmatically inspecting the Relax IR structure instead of relying on string matching.

Comment on lines +3492 to +3494

		tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)

		assert 'R.const(False, "bool")' in str(tvm_model)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test can be made more robust by inspecting the Relax IR programmatically instead of using a string search. This ensures that the structure is exactly as expected and is not sensitive to formatting changes in the string representation of the IR.

	tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)

	assert 'R.const(False, "bool")' in str(tvm_model)
	tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)
	main_func = tvm_model["main"]
	assert isinstance(main_func.body, relax.Block)
	assert len(main_func.body.bindings) == 1
	binding = main_func.body.bindings[0]
	assert isinstance(binding, relax.VarBinding)
	assert isinstance(binding.value, relax.Constant)
	assert not binding.value.data.numpy()
	assert binding.value.struct_info.dtype == "bool"

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed. I replaced the string-based IR assertion with programmatic Relax IR inspection so the test now checks the expected constant/binding structure directly and is no longer sensitive to textual IR formatting changes.

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some issues found:

Please add a guard in from_onnx for empty Optional graph outputs. Right now, if a graph output is an empty Optional, outputs can contain an _EmptyOptional object instead of a relax.Expr, and emit_output will fail with an opaque internal error. It would be better to raise a clear user-facing error here.
The _is_empty_optional branch for the "Optional" op looks dead right now, since "Optional" is already handled earlier via return_tuple_ops. Please consider keeping only one mechanism here — either remove "Optional" from return_tuple_ops and let _is_empty_optional handle it, or remove the dead branch.

Addressed, thanks for the suggestions.

I added an explicit guard in from_onnx for graph outputs: if any output resolves to an empty Optional sentinel (_EmptyOptional), the importer now raises a clear user-facing ValueError instead of letting it flow into emit_output and fail with an opaque internal error.

I also removed the redundant _is_empty_optional(op) output-count branch in _construct_nodes. Since Optional is already handled through return_tuple_ops, keeping both mechanisms was unnecessary. This keeps the output handling path single and easier to reason about.

Additionally, I added a regression test (test_empty_optional_graph_output_raises) to ensure this behavior is covered.

cc @tlopex