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 d01d44139e [ONNX] Preserve integer Div constant-fold precision (#20324)
d01d44139e is described below
commit d01d44139e2bf82ad538c60a1ab5dd5cd40f9d4e
Author: Nanmur <[email protected]>
AuthorDate: Mon Sep 14 12:04:46 2026 +0800
[ONNX] Preserve integer Div constant-fold precision (#20324)
The Relax ONNX importer constant-folds binary operations through NumPy.
For integer `Div`, `numpy.divide` promotes the operands to `float64`, so
`int64` values above `2**53` lose low bits before the result is cast
back to the integer dtype.
This change handles integer constant operands directly with integer
quotient/remainder arithmetic. NumPy floor division is adjusted by one
when signed operands have opposite signs and a non-zero remainder,
preserving ONNX's truncation-toward-zero behavior without passing
through floating point. Existing tensor and `PrimExpr` paths are
unchanged.
The regression test covers values immediately above `2**53`, a negative
large integer, and both signed truncation directions.
Fixes #20281
Tests:
- `python -m pytest tests/python/relax/test_frontend_onnx.py -k
'test_binary or test_div_integer' -q` (`13 passed`)
- `python -m ruff check python/tvm/relax/frontend/onnx/onnx_frontend.py
tests/python/relax/test_frontend_onnx.py`
- `python -m ruff format --check
python/tvm/relax/frontend/onnx/onnx_frontend.py
tests/python/relax/test_frontend_onnx.py`
---
python/tvm/relax/frontend/onnx/onnx_frontend.py | 15 +++++++++++++
tests/python/relax/test_frontend_onnx.py | 30 +++++++++++++++++++++++++
2 files changed, 45 insertions(+)
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py
b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index 62cde1ee12..498d0070fe 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -674,6 +674,15 @@ class Div(BinaryBase):
return int(expr.value) == 0
return False
+ @staticmethod
+ def _numpy_integer_divide(lhs, rhs, signed):
+ quotient, remainder = _np.divmod(lhs, rhs)
+ if signed:
+ signs_differ = _np.signbit(lhs) != _np.signbit(rhs)
+ adjust_toward_zero = _np.logical_and(signs_differ, remainder != 0)
+ quotient = quotient + adjust_toward_zero.astype(quotient.dtype)
+ return quotient
+
@classmethod
def _impl_v7(cls, bb, inputs, attr, params):
try:
@@ -700,6 +709,12 @@ class Div(BinaryBase):
if cls._is_zero(inputs[1]):
raise ValueError("ONNX Div with integer inputs encountered divisor
value 0.")
+ if all(isinstance(inp, relax.Constant) for inp in inputs):
+ lhs = inputs[0].data.numpy()
+ rhs = inputs[1].data.numpy()
+ output = cls._numpy_integer_divide(lhs, rhs, lhs_code ==
DataTypeCode.INT)
+ return relax.const(output, lhs_dtype)
+
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)
diff --git a/tests/python/relax/test_frontend_onnx.py
b/tests/python/relax/test_frontend_onnx.py
index 92f62d32f6..660179f3b4 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -685,6 +685,36 @@ def
test_div_integer_constant_folding_truncates_toward_zero():
tvm.ir.assert_structural_equal(tvm_model, Expected)
+def test_div_integer_constant_folding_preserves_int64_precision():
+ dividend_values = np.array([2**53 + 1, 2**53 + 3, -(2**53 + 3), -5, 5],
dtype=np.int64)
+ divisor_values = np.array([1, 1, 1, 2, -2], dtype=np.int64)
+ expected = np.array([2**53 + 1, 2**53 + 3, -(2**53 + 3), -2, -2],
dtype=np.int64)
+
+ a = numpy_helper.from_array(dividend_values, name="a")
+ b = numpy_helper.from_array(divisor_values, name="b")
+ node = helper.make_node("Div", ["a", "b"], ["y"])
+ graph = helper.make_graph(
+ [node],
+ "div_integer_constant_precision",
+ [],
+ [helper.make_tensor_value_info("y", TensorProto.INT64, [5])],
+ initializer=[a, b],
+ )
+ model = helper.make_model(graph, opset_imports=[helper.make_opsetid("",
18)])
+ model.ir_version = 9
+
+ tvm_model = from_onnx(model, opset=18, keep_params_in_input=False)
+ folded_outputs = []
+
+ def collect_constants(expr):
+ if isinstance(expr, relax.Constant):
+ folded_outputs.append(expr.data.numpy())
+
+ relax.analysis.post_order_visit(tvm_model["main"].body, collect_constants)
+ assert len(folded_outputs) == 1
+ np.testing.assert_array_equal(folded_outputs[0], expected)
+
+
@pytest.mark.parametrize(
("input_size", "divisor_shape", "offset"),
[(386, [], None), (384, [1], 2)],