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 22ee81e569 [Fix][Relax][ONNX] Preserve integer Div truncation during
import (#19975)
22ee81e569 is described below
commit 22ee81e56917cb35685bf41dd1cf30146d2720ee
Author: Vic Wen <[email protected]>
AuthorDate: Wed Jul 15 07:36:41 2026 +0800
[Fix][Relax][ONNX] Preserve integer Div truncation during import (#19975)
ONNX integer Div uses truncating division, rounding toward zero. The
Relax ONNX frontend already special-cased integer Div to detect zero
divisors, but its PrimExpr folding path could still use NumPy
floating-point division when one of the inputs was a shape-derived
PrimExpr.
That behavior can produce floating-point TIR values for integer
shape/index computations. For example, a `Shape -> Gather -> Div ->
Slice` subgraph can produce `T.float64(128.666...)` as a Slice bound,
which Relax rejects because strided_slice expects integer PrimExpr
bounds.
This patch handles scalar integer Div inputs that contain a PrimExpr
using TIR `truncdiv`, preserving ONNX semantics while keeping shape
computations in TIR instead of routing them through NumPy. Constant
tensor Div continues to use the existing generic binary constant-folding
path.
The regression tests cover:
- integer constant folding with negative values to distinguish
truncation from floor division
- a shape-derived PrimExpr Div used as a Slice bound
- integer zero-divisor error handling
Verification:
- `python -m pytest
tests/python/relax/test_frontend_onnx.py::test_div_integer_constant_zero_divisor_raises_valueerror
tests/python/relax/test_frontend_onnx.py::test_div_integer_constant_folding_truncates_toward_zero
tests/python/relax/test_frontend_onnx.py::test_div_integer_primexpr_folding_truncates_toward_zero
-q`
Fixes #19974
Signed-off-by: viiccwen <[email protected]>
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 40 ++++++++++++-
tests/python/relax/test_frontend_onnx.py | 80 +++++++++++++++++++++++++
2 files changed, 117 insertions(+), 3 deletions(-)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 7f904c8a1b..f195f6276a 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -640,11 +640,39 @@ class Div(BinaryBase):
numpy_op = _np.divide
relax_op = relax.op.divide
+ @staticmethod
+ def _as_scalar_prim_expr(expr, dtype):
+ if tvm.ir.is_prim_expr(expr):
+ return expr
+ if isinstance(expr, relax.Constant):
+ data = expr.data.numpy()
+ if data.size == 1:
+ return tirx.const(data.item(), dtype)
+ return None
+
+ @staticmethod
+ def _is_zero(expr):
+ if isinstance(expr, relax.Constant):
+ return bool(_np.any(expr.data.numpy() == 0))
+ if isinstance(expr, tirx.IntImm):
+ return int(expr.value) == 0
+ return False
+
@classmethod
def _impl_v7(cls, bb, inputs, attr, params):
try:
- lhs_code = DataType(inputs[0].ty.dtype.dtype).type_code
- rhs_code = DataType(inputs[1].ty.dtype.dtype).type_code
+ lhs_dtype = (
+ str(getattr(inputs[0], "dtype", None) or inputs[0].ty)
+ if tvm.ir.is_prim_expr(inputs[0])
+ else inputs[0].ty.dtype.dtype
+ )
+ rhs_dtype = (
+ str(getattr(inputs[1], "dtype", None) or inputs[1].ty)
+ if tvm.ir.is_prim_expr(inputs[1])
+ else inputs[1].ty.dtype.dtype
+ )
+ lhs_code = DataType(lhs_dtype).type_code
+ rhs_code = DataType(rhs_dtype).type_code
except (AttributeError, ValueError, TypeError, RuntimeError):
return cls.base_impl(bb, inputs, attr, params)
@@ -653,9 +681,15 @@ class Div(BinaryBase):
if not (lhs_is_integer and rhs_is_integer):
return cls.base_impl(bb, inputs, attr, params)
- if isinstance(inputs[1], relax.Constant) and
bool(_np.any(inputs[1].data.numpy() == 0)):
+ if cls._is_zero(inputs[1]):
raise ValueError("ONNX Div with integer inputs encountered divisor
value 0.")
+ has_prim_expr = any(tvm.ir.is_prim_expr(inp) for inp in inputs)
+ lhs = cls._as_scalar_prim_expr(inputs[0], lhs_dtype)
+ rhs = cls._as_scalar_prim_expr(inputs[1], rhs_dtype)
+ if has_prim_expr and lhs is not None and rhs is not None:
+ return relax.prim_value(tirx.truncdiv(lhs, rhs))
+
return cls.base_impl(bb, inputs, attr, params)
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 30d9c35a6c..8a06b5142d 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -669,6 +669,86 @@ def
test_div_integer_constant_zero_divisor_raises_valueerror():
from_onnx(model, opset=18, keep_params_in_input=False)
+def test_div_integer_constant_folding_truncates_toward_zero():
+ a = make_constant_node("a", TensorProto.INT64, [2], [-5, 5])
+ b = make_constant_node("b", TensorProto.INT64, [2], [2, 2])
+ node = helper.make_node("Div", ["a", "b"], ["y"])
+ graph = helper.make_graph(
+ [a, b, node],
+ "div_integer_constant",
+ [],
+ [helper.make_tensor_value_info("y", TensorProto.INT64, [2])],
+ )
+ model = helper.make_model(graph, opset_imports=[helper.make_opsetid("",
13)])
+ model.ir_version = 8
+
+ tvm_model = from_onnx(model, opset=13)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main() -> R.Tensor((2,), dtype="int64"):
+ R.func_attr({"num_input": 0})
+ with R.dataflow():
+ gv: R.Tensor((2,), dtype="int64") = R.const([-2, 2], "int64")
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(tvm_model, Expected)
+
+
[email protected](
+ ("input_size", "divisor_shape", "offset"),
+ [(386, [], None), (384, [1], 2)],
+)
+def test_div_integer_primexpr_folding_truncates_toward_zero(input_size,
divisor_shape, offset):
+ shape = helper.make_node("Shape", ["x"], ["x_shape"])
+ axis = make_constant_node("axis", TensorProto.INT64, [], [0])
+ dim = helper.make_node("Gather", ["x_shape", "axis"], ["dim"])
+ nodes = [shape, axis, dim]
+ dividend = "dim"
+ if offset is not None:
+ offset_node = make_constant_node("offset", TensorProto.INT64, [],
[offset])
+ shifted_dim = helper.make_node("Add", ["dim", "offset"],
["shifted_dim"])
+ nodes.extend([offset_node, shifted_dim])
+ dividend = "shifted_dim"
+
+ divisor = make_constant_node("divisor", TensorProto.INT64, divisor_shape,
[3])
+ end = helper.make_node("Div", [dividend, "divisor"], ["end"])
+ starts = make_constant_node("starts", TensorProto.INT64, [1], [0])
+ axes = make_constant_node("axes", TensorProto.INT64, [1], [0])
+ steps = make_constant_node("steps", TensorProto.INT64, [1], [1])
+ slice_node = helper.make_node("Slice", ["x", "starts", "end", "axes",
"steps"], ["y"])
+ nodes.extend([divisor, end, starts, axes, steps, slice_node])
+
+ graph = helper.make_graph(
+ nodes,
+ "div_integer_primexpr",
+ [helper.make_tensor_value_info("x", TensorProto.FLOAT, [input_size])],
+ [helper.make_tensor_value_info("y", TensorProto.FLOAT, [128])],
+ )
+ model = helper.make_model(graph, opset_imports=[helper.make_opsetid("",
13)])
+ model.ir_version = 8
+
+ tvm_model = from_onnx(model, opset=13)
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(
+ x: R.Tensor((input_size,), dtype="float32"),
+ ) -> R.Tensor((128,), dtype="float32"):
+ R.func_attr({"num_input": 1})
+ with R.dataflow():
+ gv: R.Tensor((128,), dtype="float32") = R.strided_slice(
+ x, axes=[0], begin=[0], end=[128], strides=[1],
assume_inbound=False
+ )
+ R.output(gv)
+ return gv
+
+ tvm.ir.assert_structural_equal(tvm_model, Expected)
+
+
@pytest.mark.parametrize("int_mode", [True, False])
def test_mod(int_mode: bool):
if int_mode:
