diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py index 929e2e278138..76951718c497 100644 --- a/python/tvm/relax/backend/dispatch_sort_scan.py +++ b/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 diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py index e9b951e0f9f5..576214293d1e 100644 --- a/python/tvm/relax/frontend/onnx/onnx_frontend.py +++ b/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) diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index f96984e42375..25f5b302b09e 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/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,