Aharrypotter opened a new pull request, #20126:
URL: https://github.com/apache/tvm/pull/20126
## Summary
This PR extends Relax `quantize` and `dequantize` type inference to support
rank-zero inputs when their scale and zero point describe per-tensor
quantization.
The gap was exposed by the INT8 dynamic-shape BiRefNeXt anime segmentation
model. Its opset-18 graph contains scalar-input `DequantizeLinear` nodes with
scalar scale and zero-point initializers. This is valid ONNX per-tensor
dequantization, but the current Relax operator rejects the canonicalized axis
because a rank-zero tensor has no axis.
## Goal
Import and execute valid scalar ONNX `QuantizeLinear` and `DequantizeLinear`
graphs without weakening validation for per-axis quantization.
## What changed
- Treat the axis as irrelevant for rank-zero Relax quantize/dequantize
inputs.
- Require rank-zero inputs to use scalar or singleton scale and zero-point
tensors.
- Keep the existing axis and parameter-size checks unchanged for inputs with
rank one or greater.
- Add symmetric core type-inference coverage for scalar quantize and
dequantize.
- Add ONNX Runtime-backed numerical and dtype coverage for scalar opset-18
`QuantizeLinear` and `DequantizeLinear` graphs.
## Design
ONNX uses scalar scale and zero-point tensors for per-tensor quantization. In
that mode the `axis` attribute is ignored. The Relax ONNX frontend already
canonicalizes the default axis to `0` for inputs with rank at most one, but
the
Relax operator type relation previously rejected every axis for a rank-zero
input:
```text
rank(input) = 0
valid axis range = [0, rank(input) - 1] = [0, -1]
```
For rank-zero input, this change therefore skips the axis-range check and
instead validates the condition that makes the axis irrelevant: both scale
and
zero point must be scalar or singleton tensors. The legalization and output
shape rules already handle rank-zero tensors, so no converter or lowering
special case is needed.
## Updated behavior
| Input / quantization parameters | Behavior |
| --- | --- |
| Rank zero + scalar/singleton scale and zero point | Accept as per-tensor
QDQ; preserve scalar output shape |
| Rank zero + non-singleton scale or zero point | Reject with an explicit
error |
| Rank one or greater | Preserve the existing axis and parameter-size
validation |
## Safety checks
- Scalar support is implemented symmetrically for quantize and dequantize.
- A parametrized core test verifies that non-singleton per-axis parameters
are
still rejected for rank-zero input.
- The ONNX tests compare TVM execution and output dtypes with ONNX Runtime.
- The full Relax ONNX frontend test file passes in the validation
environment,
apart from five pre-existing Float8 saturate baseline cases that were
excluded explicitly.
## Out of scope / non-goals
- Claiming full BiRefNeXt import or execution; this PR removes only its first
observed scalar-QDQ blocker, and the model still has an independent dynamic
shape `Concat` blocker.
- Adding the external 117 MB BiRefNeXt model to the TVM test suite.
- Changing per-axis QDQ behavior for non-scalar inputs.
- Addressing the existing Float8 saturate test baseline.
## Results
Before this change, minimal scalar ONNX graphs fail during import with:
```text
ValueError: relax.quantize: axis param is out of range (0)
ValueError: relax.dequantize: axis param is out of range (0)
```
With this change, opset-18 scalar graphs import, legalize, compile, and
execute
against ONNX Runtime:
```text
QuantizeLinear: x = 1.25, scale = 0.25, zero_point = 2 -> uint8 scalar 7
DequantizeLinear: x = 7, scale = 0.25, zero_point = 2 -> float32 scalar
1.25
```
## Tests
- `pre-commit run --files src/relax/op/tensor/qdq.cc
tests/python/relax/test_frontend_onnx.py tests/python/relax/test_op_qdq.py`
- Fresh CPU-only H20 Release build: passed
- Relax QDQ operator tests: `9 passed`
- Focused scalar ONNX Runtime-backed tests: `2 passed, 500 deselected`
- Relax ONNX frontend H20 run: `484 passed, 9 skipped, 5 deselected, 4
xfailed`
## References
- [ONNX QuantizeLinear
specification](https://onnx.ai/onnx/operators/onnx__QuantizeLinear.html)
- [ONNX DequantizeLinear
specification](https://onnx.ai/onnx/operators/onnx__DequantizeLinear.html)
- [BiRefNeXt anime segmentation ONNX
model](https://huggingface.co/nkta/birefnext-aniseg-ONNX)
--
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]
