[Fix][Relax][ONNX] Preserve integer Div truncation during import

python/tvm/relax/backend/dispatch_sort_scan.py

@@ -86,9 +86,7 @@ def visit_call_(self, call: relax.Call) -> relax.Expr:
                 if self.is_gpu_target(tgt):
                     te_func = topi.gpu.searchsorted
             out_dtype = "int32" if call.attrs.out_int32 else "int64"
-            return self.builder_.call_te(
-                te_func, boundaries, input_tensor, right, out_dtype
-            )
+            return self.builder_.call_te(te_func, boundaries, input_tensor, right, out_dtype)
         if call.op.name == "relax.sort":
             tgt = self._get_target(call.ty)
             te_func = topi.sort

python/tvm/relax/frontend/onnx/onnx_frontend.py

@@ -527,11 +527,37 @@ 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) and expr.data.shape == ():
+            return tirx.const(expr.data.numpy().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)
 
@@ -540,9 +566,15 @@ def _impl_v7(cls, bb, inputs, attr, params):
         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)
 
 

tests/python/relax/test_frontend_onnx.py

@@ -669,6 +669,70 @@ 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)
+
+
+def test_div_integer_primexpr_folding_truncates_toward_zero():
+    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"])
+    divisor = make_constant_node("divisor", TensorProto.INT64, [], [3])
+    end = helper.make_node("Div", ["dim", "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"])
+    graph = helper.make_graph(
+        [shape, axis, dim, divisor, end, starts, axes, steps, slice_node],
+        "div_integer_primexpr",
+        [helper.make_tensor_value_info("x", TensorProto.FLOAT, [386])],
+        [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((386,), 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:
@@ -6563,7 +6627,9 @@ def verify_pad(input_shape, pads, expected, mode="constant", value=0.0, opset=14
             if axes is not None:
                 axes = np.array(axes, dtype=np.int64)
                 node_inputs = ["input", "pads", "", "axes"]
-                initializer.append(helper.make_tensor("axes", TensorProto.INT64, (len(axes),), axes))
+                initializer.append(
+                    helper.make_tensor("axes", TensorProto.INT64, (len(axes),), axes)
+                )
 
             node = helper.make_node("Pad", inputs=node_inputs, outputs=["output"], mode=mode)
             graph = helper.make_graph(
@@ -6628,7 +6694,9 @@ def verify_pad(input_shape, pads, expected, mode="constant", value=0.0, opset=14
         verify_pad(
             input_shape,
             pads,
-            _make_pad_expected_ir(input_shape, pads, mode=mode, value=value, opset=opset, axes=axes),
+            _make_pad_expected_ir(
+                input_shape, pads, mode=mode, value=value, opset=opset, axes=axes
+            ),
             mode,
             value,
             opset,