This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new b01df356ef [Relax][Frontend][ONNX] Support symbolic shapes in Min/Max
broadcast (#20218)
b01df356ef is described below
commit b01df356ef5c782be296c87be43aba4055fad4d2
Author: Neo Chien <[email protected]>
AuthorDate: Sun Aug 30 09:00:06 2026 +0800
[Relax][Frontend][ONNX] Support symbolic shapes in Min/Max broadcast
(#20218)
Hi Committers,
This PR fixes issue https://github.com/apache/tvm/issues/20175. Any
suggestions would be appreciated if you are available.
### Root Cause
The `compute_broadcast_shape()` function in the ONNX frontend attempts
to use the built-in `max()` function and boolean operators to broadcast
dimensions. These operations fail when the dimensions are symbolic TIR
expressions because:
- `Expr` objects cannot be used in boolean conditions.
- Built-in `max()` requires comparison operators that function within a
boolean context.
- The original implementation assumed all dimensions were of the same
type.
### Solution
Implement three new components to handle dimension-type-aware
broadcasting:
- `_normalize_shape_dim()` converts IntImm constants to ints, ensuring
consistent type handling.
- `_broadcast_shape_dims()` implements broadcast rules
- Refactor `compute_broadcast_shape()` to use `_broadcast_shape_dims()`
---------
Co-authored-by: cchung100m <[email protected]>
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 66 +++++++-
tests/python/relax/test_frontend_onnx.py | 215 ++++++++++++++++++++++++
2 files changed, 277 insertions(+), 4 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 8a40e4a612..1afbf208b2 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -58,6 +58,7 @@ import tvm_ffi
import tvm
from tvm import relax, tirx, topi
+from tvm.ir import Expr as TvmExpr
from tvm.ir import IRModule
from tvm.ir.supply import UniqueNameSupply
from tvm.runtime import DataType, DataTypeCode
@@ -2618,6 +2619,66 @@ class Sqrt(OnnxOpConverter):
return relax.op.sqrt(inputs[0])
+def _normalize_shape_dim(dim: int | tirx.Expr) -> int | tirx.Expr:
+ """Normalize a shape dimension, converting IntImm to int if needed."""
+ if isinstance(dim, tirx.IntImm):
+ return int(dim.value)
+ return dim
+
+
+def _broadcast_shape_dims(ai: int | tirx.Expr, bi: int | tirx.Expr) -> int |
tirx.Expr:
+ """Broadcast two shape dimensions according to ONNX broadcasting rules.
+
+ Handles both constant and symbolic dimensions. Returns the broadcasted
dimension.
+ Raises ValueError if the dimensions are incompatible.
+
+ Broadcasting rules:
+ - If one dimension is 1, broadcast to the other dimension
+ - If both dimensions are equal (same value or same symbolic variable), use
that dimension
+ - Otherwise, dimensions are incompatible
+ """
+ ai = _normalize_shape_dim(ai)
+ bi = _normalize_shape_dim(bi)
+
+ # Check if either dimension is symbolic (Expr)
+ ai_is_expr = isinstance(ai, TvmExpr)
+ bi_is_expr = isinstance(bi, TvmExpr)
+
+ if not ai_is_expr and not bi_is_expr:
+ # Both are constants: unchanged from original
+ if ai == bi or ai == 1 or bi == 1:
+ return max(ai, bi)
+ else:
+ raise ValueError(f"Cannot broadcast {ai} and {bi}")
+ elif ai_is_expr and bi_is_expr:
+ # Both are symbolic: Check if they're the same variable
+ # If so, return one of them; otherwise, cannot determine at import time
+ is_equal = tvm_ffi.structural_equal(ai, bi)
+
+ if is_equal:
+ return ai
+ else:
+ raise ValueError(
+ f"Cannot broadcast symbolic dimensions {ai} and {bi} - "
+ "both are symbolic but structurally different. "
+ "ONNX frontend requires dimensions names to match for symbolic
shapes."
+ )
+ else:
+ # One is symbolic, one is constant
+ const_val = bi if ai_is_expr else ai
+ expr_val = ai if ai_is_expr else bi
+
+ if const_val == 1:
+ # Constant dimension is 1, broadcast to symbolic dimension
+ return expr_val
+ else:
+ # Constant dimension is not 1
+ raise ValueError(
+ f"Cannot broadcast symbolic dimension {expr_val} with non-1
constant {const_val}: "
+ f"runtime value of symbolic dimension is unknown at compile
time."
+ )
+
+
def compute_broadcast_shape(shape_a, shape_b):
"""Compute target shape for Multidirectional Broadcasting"""
rank = max(len(shape_a), len(shape_b))
@@ -2627,10 +2688,7 @@ def compute_broadcast_shape(shape_a, shape_b):
target = []
for ai, bi in zip(a, b):
- if ai == bi or ai == 1 or bi == 1:
- target.append(max(ai, bi))
- else:
- raise ValueError(f"Cannot broadcast {ai} and {bi}")
+ target.append(_broadcast_shape_dims(ai, bi))
return tuple(target)
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index a2081adab9..7fb80d55dc 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -928,6 +928,221 @@ def test_multi_input_constant_rank_axis_bounds(op_name,
rank):
check_correctness(helper.make_model(graph), opset=13)
+def _make_onnx_min_max_model(
+ x_shape: list,
+ y_shape: list,
+ out_shape: list,
+ op_name: str = "Min",
+ graph_name: str = "test_graph",
+) -> onnx.ModelProto:
+ """Construct a Min/Max ONNX model.
+
+ Args:
+ x_shape: Shape of the first input tensor
+ y_shape: Shape of the second input tensor
+ out_shape: Shape of the output tensor
+ op_name: The operation name, either "Min" or "Max"
+ graph_name: The name of the ONNX graph
+
+ Returns:
+ An ONNX ModelProto object representing the Min/Max operation.
+ """
+ x = helper.make_tensor_value_info("x", TensorProto.FLOAT, x_shape)
+ y = helper.make_tensor_value_info("y", TensorProto.FLOAT, y_shape)
+ out = helper.make_tensor_value_info("out", TensorProto.FLOAT, out_shape)
+
+ graph = helper.make_graph(
+ [helper.make_node(op_name, ["x", "y"], ["out"])],
+ graph_name,
+ [x, y],
+ [out],
+ )
+ return helper.make_model(graph, opset_imports=[helper.make_opsetid("",
18)])
+
+
+def _make_expected_broadcast_ir_min(
+ x_shape: tuple,
+ y_shape: tuple,
+):
+ """Generate expected broadcast IR module for Min operation.
+
+ Args:
+ x_shape: Shape of the first input tensor
+ y_shape: Shape of the second input tensor
+
+ Returns:
+ Expected IR module for the Min operation.
+ """
+ output_shape = (x_shape[0], 4)
+
+ @I.ir_module
+ class ExpectedMin:
+ @R.function
+ def main(
+ x: R.Tensor(x_shape, dtype="float32"),
+ y: R.Tensor(y_shape, dtype="float32"),
+ ) -> R.Tensor(output_shape, dtype="float32"):
+ n = T.int64()
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ lv = R.broadcast_to(x, R.shape((n, 4)))
+ lv1 = R.broadcast_to(y, R.shape((n, 4)))
+ lv2 = R.stack((lv, lv1), axis=0)
+ gv = R.min(lv2, axis=[0], keepdims=False)
+ R.output(gv)
+ return gv
+
+ return ExpectedMin
+
+
+def _make_expected_broadcast_ir_max(
+ x_shape: tuple,
+ y_shape: tuple,
+):
+ """Generate expected broadcast IR module for Max operation.
+
+ Args:
+ x_shape: Shape of the first input tensor
+ y_shape: Shape of the second input tensor
+
+ Returns:
+ Expected IR module for the Max operation.
+ """
+ output_shape = (x_shape[0], 4)
+
+ @I.ir_module
+ class ExpectedMax:
+ @R.function
+ def main(
+ x: R.Tensor(x_shape, dtype="float32"),
+ y: R.Tensor(y_shape, dtype="float32"),
+ ) -> R.Tensor(output_shape, dtype="float32"):
+ n = T.int64()
+ R.func_attr({"num_input": 2})
+ with R.dataflow():
+ lv = R.broadcast_to(x, R.shape((n, 4)))
+ lv1 = R.broadcast_to(y, R.shape((n, 4)))
+ lv2 = R.stack((lv, lv1), axis=0)
+ gv = R.max(lv2, axis=[0], keepdims=False)
+ R.output(gv)
+ return gv
+
+ return ExpectedMax
+
+
+def _test_symbolic_broadcast_case(
+ x_shape: list,
+ y_shape: list,
+ out_shape: list,
+ op_name: str,
+ graph_name: str,
+ should_pass: bool = True,
+ error_pattern: str | None = None,
+):
+ """Execute a single symbolic broadcast test case.
+
+ Args:
+ x_shape: Shape of the first input tensor
+ y_shape: Shape of the second input tensor
+ out_shape: Shape of the output tensor
+ op_name: The operation name, either "Min" or "Max"
+ graph_name: The name of the ONNX graph
+ should_pass: Whether the test is expected to pass or raise an error
+ error_pattern: The expected error message pattern if should_pass is
False
+ """
+ model = _make_onnx_min_max_model(x_shape, y_shape, out_shape, op_name,
graph_name)
+ onnx.checker.check_model(model)
+
+ if should_pass:
+ tvm_model = from_onnx(model, opset=18, keep_params_in_input=True)
+ if op_name == "Min":
+ expected = _make_expected_broadcast_ir_min(tuple(x_shape),
tuple(y_shape))
+ else:
+ expected = _make_expected_broadcast_ir_max(tuple(x_shape),
tuple(y_shape))
+ tvm.ir.assert_structural_equal(tvm_model, expected)
+ else:
+ with pytest.raises(ValueError, match=error_pattern):
+ from_onnx(model, opset=18, keep_params_in_input=True)
+
+
+def test_multi_input_broadcasting_symbolic_shapes():
+ """Comprehensive test for symbolic shape broadcasting in Min/Max
operations.
+
+ Tests four levels:
+ 1. LEVEL 1 - Compile-time: Import and symbolic dimension handling (4 cases)
+ 2. LEVEL 2 - Shape lowering: Pass through LegalizeOps pipeline
+ 3. LEVEL 3 - IR verification: Check broadcast_to operations
+ 4. LEVEL 4 - Conservative rejection: Verify symbolic vs non-1 constant is
rejected
+
+ When dimensions are symbolic (not constant), broadcasting should handle
them correctly
+ without raising "Cannot use and / or / not operator to Expr".
+ Conservative rejection prevents silent miscompilation where TOPI cannot
correctly handle
+ symbolic source dimensions in broadcast operations.
+ """
+ # LEVEL 1 - Compile-time: Import and symbolic dimension handling (4 cases)
+ _test_symbolic_broadcast_case(
+ x_shape=["n", 4],
+ y_shape=["n", 4],
+ out_shape=["n", 4],
+ op_name="Min",
+ graph_name="min_symbolic_broadcast",
+ )
+
+ _test_symbolic_broadcast_case(
+ x_shape=["n", 4],
+ y_shape=["n", 4],
+ out_shape=["n", 4],
+ op_name="Max",
+ graph_name="max_symbolic_broadcast",
+ )
+
+ _test_symbolic_broadcast_case(
+ x_shape=["n", 1],
+ y_shape=["n", 4],
+ out_shape=["n", 4],
+ op_name="Min",
+ graph_name="min_mixed_symbolic_broadcast",
+ )
+
+ _test_symbolic_broadcast_case(
+ x_shape=["n", 4],
+ y_shape=["m", 4],
+ out_shape=["n", 4],
+ op_name="Min",
+ graph_name="min_different_symbolic",
+ should_pass=False,
+ error_pattern="Cannot broadcast symbolic dimensions",
+ )
+
+ # LEVEL 2 - Shape lowering: Pass through LegalizeOps pipeline
+ model_lower = _make_onnx_min_max_model(
+ ["n", 4], ["n", 4], ["n", 4], "Min", "min_shape_lower_test"
+ )
+ onnx.checker.check_model(model_lower)
+
+ tvm_model = from_onnx(model_lower, opset=18, keep_params_in_input=True)
+
+ tvm_model = relax.transform.DecomposeOpsForInference()(tvm_model)
+ tvm_model = relax.transform.LegalizeOps()(tvm_model)
+
+ tvm_model, _ = relax.frontend.detach_params(tvm_model)
+
+ # LEVEL 3 - IR verification: Check broadcast_to operations
+ ir_str = str(tvm_model)
+ assert "broadcast" in ir_str, "IR should contain broadcast operation"
+
+ # LEVEL 4 - Conservative rejection: Verify symbolic vs non-1 constant is
rejected
+ _test_symbolic_broadcast_case(
+ x_shape=["n", 4],
+ y_shape=[2, 4],
+ out_shape=[2, 4],
+ op_name="Min",
+ graph_name="min_symbolic_vs_constant",
+ should_pass=False,
+ error_pattern="runtime value of symbolic dimension is unknown",
+ )
+
+
@pytest.mark.parametrize("op_name", ["Min", "Max", "Sum", "Mean"])
@pytest.mark.parametrize("shape", [[], [5], [2, 3]])
def test_multi_input_constant_single_input(op_name, shape):
