diff --git a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py index 3ff3b596af3b..9c49abb913b4 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -246,6 +246,38 @@ def _emit_torch_reshape(self, x, dims): x = self.block_builder.emit(relax.op.reshape(x, step)) return x + @staticmethod + def _scalar_result_dtype(tensor_dtype, scalar) -> str | None: + """Return the dtype torch gives ``tensor scalar`` for a Python scalar operand. + + A Python scalar takes part in type promotion at a lower priority than a tensor: it + widens the tensor only when it belongs to a higher category. So a float scalar + promotes an integer or bool tensor to the default float dtype, an int scalar + promotes only a bool tensor (to int64), and otherwise the tensor's dtype wins. + Casting the scalar down to the tensor's dtype instead turns ``x * 0.5`` on an + integer tensor into ``x * 0``. Returns None for a dtype torch cannot map. + """ + import torch # type: ignore + + if not isinstance(scalar, bool | int | float): + return None + torch_dtype = { + "float64": torch.float64, + "float32": torch.float32, + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "int64": torch.int64, + "int32": torch.int32, + "int16": torch.int16, + "int8": torch.int8, + "uint8": torch.uint8, + "bool": torch.bool, + }.get(str(tensor_dtype)) + if torch_dtype is None: + return None + promoted = torch.result_type(torch.empty(0, dtype=torch_dtype), scalar) + return str(promoted).replace("torch.", "") + @staticmethod def _promote_common_dtype(lhs_dtype: str | None, rhs_dtype: str | None) -> str | None: """Return the promoted dtype following PyTorch rules, or None if unsupported.""" @@ -624,6 +656,13 @@ def _round(self, node: fx.Node) -> relax.Expr: result = relax.op.astype(result, input_dtype) return self.block_builder.emit(result) + def _reciprocal(self, node: fx.Node) -> relax.Var: + # torch.reciprocal is 1 / x under true-division rules: an integer or bool input + # comes back as the default float dtype, not as an integer quotient. + x = self.env[node.args[0]] + one, x = self._true_division_operands(1, x) + return self.block_builder.emit(relax.op.divide(one, x)) + def _softmax(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("dim", -1) @@ -687,35 +726,63 @@ def convert(node: fx.Node) -> relax.Var: ########## Binary Ops ########## + def _promote_scalar_operand(self, tensor, scalar): + """Widen ``tensor`` if a Python ``scalar`` outranks it, and build the constant. + + torch.result_type decides who wins: the tensor is only widened when the scalar's + category is higher (float scalar vs int tensor, int scalar vs bool tensor). The + constant is then built in that dtype rather than truncated to the tensor's. + """ + target = self._scalar_result_dtype(tensor.ty.dtype, scalar) or tensor.ty.dtype + if str(tensor.ty.dtype) != str(target): + tensor = self.block_builder.emit(relax.op.astype(tensor, target)) + return tensor, relax.const(scalar, target) + + def _promote_binary_operands(self, lhs, rhs): + """Bring the two operands of a binary op to torch's promoted dtype.""" + if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr): + lhs_si = getattr(lhs, "ty", None) + rhs_si = getattr(rhs, "ty", None) + if isinstance(lhs_si, relax.TensorType) and isinstance(rhs_si, relax.TensorType): + target_dtype = self._promote_common_dtype(lhs_si.dtype, rhs_si.dtype) + if target_dtype is not None: + if lhs_si.dtype != target_dtype: + lhs = self.block_builder.emit(relax.op.astype(lhs, target_dtype)) + if rhs_si.dtype != target_dtype: + rhs = self.block_builder.emit(relax.op.astype(rhs, target_dtype)) + return lhs, rhs + elif isinstance(lhs, relax.Expr): + assert isinstance(lhs.ty, relax.TensorType) + return self._promote_scalar_operand(lhs, rhs) + elif isinstance(rhs, relax.Expr): + assert isinstance(rhs.ty, relax.TensorType) + rhs, lhs = self._promote_scalar_operand(rhs, lhs) + return lhs, rhs + else: + assert False + + def _true_division_operands(self, lhs, rhs): + """Promote for ``a / b``: integer and bool operands divide as the default float. + + torch's true division always produces a floating result -- ``int64 / int64`` and + ``int64 / 2`` are float32, not a truncating integer quotient -- so after the usual + promotion an integral or bool pair is cast to the default float dtype. + """ + lhs, rhs = self._promote_binary_operands(lhs, rhs) + dtype = getattr(getattr(lhs, "ty", None), "dtype", None) + if dtype is not None and ( + dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT) or str(dtype) == "bool" + ): + lhs = self.block_builder.emit(relax.op.astype(lhs, "float32")) + rhs = self.block_builder.emit(relax.op.astype(rhs, "float32")) + return lhs, rhs + def _binary_op(self, relax_op: Callable, intrinsic_op: Callable) -> Callable: from torch import fx def convert(node: fx.Node) -> relax.Var: - def promote_binary_op_args(lhs, rhs): - if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr): - lhs_si = getattr(lhs, "ty", None) - rhs_si = getattr(rhs, "ty", None) - if isinstance(lhs_si, relax.TensorType) and isinstance( - rhs_si, relax.TensorType - ): - target_dtype = self._promote_common_dtype(lhs_si.dtype, rhs_si.dtype) - if target_dtype is not None: - if lhs_si.dtype != target_dtype: - lhs = self.block_builder.emit(relax.op.astype(lhs, target_dtype)) - if rhs_si.dtype != target_dtype: - rhs = self.block_builder.emit(relax.op.astype(rhs, target_dtype)) - return lhs, rhs - elif isinstance(lhs, relax.Expr): - assert isinstance(lhs.ty, relax.TensorType) - return lhs, relax.const(rhs, lhs.ty.dtype) - elif isinstance(rhs, relax.Expr): - assert isinstance(rhs.ty, relax.TensorType) - return relax.const(lhs, rhs.ty.dtype), rhs - else: - assert False - def call_binary_op(op, lhs, rhs): - lhs, rhs = promote_binary_op_args(lhs, rhs) + lhs, rhs = self._promote_binary_operands(lhs, rhs) return self.block_builder.emit(op(lhs, rhs)) lhs, rhs = self.retrieve_args(node) @@ -725,9 +792,9 @@ def call_binary_op(op, lhs, rhs): ): return call_binary_op(relax_op, lhs, rhs) elif isinstance(lhs, relax.expr.Constant) and not isinstance(rhs, relax.expr.Constant): - return call_binary_op(relax_op, lhs, relax.const(rhs, dtype=lhs.ty.dtype)) + return call_binary_op(relax_op, lhs, rhs) elif isinstance(rhs, relax.expr.Constant) and not isinstance(lhs, relax.expr.Constant): - return call_binary_op(relax_op, relax.const(lhs, dtype=rhs.ty.dtype), rhs) + return call_binary_op(relax_op, lhs, rhs) return intrinsic_op(lhs, rhs) return convert @@ -753,31 +820,37 @@ def _pow(self, node: fx.Node) -> relax.Var: return result return self._binary_op(relax.op.power, operator.pow)(node) + def _true_divide(self, node: fx.Node) -> relax.Var: + lhs, rhs = self.retrieve_args(node) + if not isinstance(lhs, relax.Expr) and not isinstance(rhs, relax.Expr): + return operator.truediv(lhs, rhs) + lhs, rhs = self._true_division_operands(lhs, rhs) + return self.block_builder.emit(relax.op.divide(lhs, rhs)) + def _div(self, node: fx.Node) -> relax.Var: args = self.retrieve_args(node) - inp_1 = args[0] - inp_2 = args[1] - - # Handle scalar cases - if isinstance(inp_2, int | float): - inp_2 = relax.const(inp_2) - - # Get rounding_mode from node kwargs + lhs, rhs = args[0], args[1] rounding_mode = args[2] if len(node.args) > 2 else node.kwargs.get("rounding_mode", None) - # Perform division based on rounding mode if rounding_mode is None: - # True division (normal float division) - return self.block_builder.emit(relax.op.divide(inp_1, inp_2)) - elif rounding_mode == "floor": - # Floor division - return self.block_builder.emit(relax.op.floor_divide(inp_1, inp_2)) - elif rounding_mode == "trunc": - # Trunc division: perform true division then truncate - true_div = self.block_builder.emit(relax.op.divide(inp_1, inp_2)) + lhs, rhs = self._true_division_operands(lhs, rhs) + return self.block_builder.emit(relax.op.divide(lhs, rhs)) + + # With a rounding mode the result keeps the promoted dtype: an integer pair + # stays integer. Both operands are promoted first so a Python scalar lands in + # the tensor's dtype (or widens it) instead of arriving as a dtype-less + # constant that fails the same-dtype check on every float tensor. + lhs, rhs = self._promote_binary_operands(lhs, rhs) + if rounding_mode == "floor": + return self.block_builder.emit(relax.op.floor_divide(lhs, rhs)) + if rounding_mode == "trunc": + dtype = getattr(getattr(lhs, "ty", None), "dtype", None) + if dtype is not None and dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT): + # Integer division in relax truncates toward zero already. + return self.block_builder.emit(relax.op.divide(lhs, rhs)) + true_div = self.block_builder.emit(relax.op.divide(lhs, rhs)) return self.block_builder.emit(relax.op.trunc(true_div)) - else: - raise ValueError(f"Unsupported rounding_mode: {rounding_mode}") + raise ValueError(f"Unsupported rounding_mode: {rounding_mode}") def _fmod(self, node: fx.Node): args = self.retrieve_args(node) diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 86c936723bf2..465185424bfd 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -98,10 +98,6 @@ def _log1p(self, node: fx.Node) -> relax.Var: one = relax.const(1, x.ty.dtype.dtype) return self.block_builder.emit(relax.op.log(relax.op.add(x, one))) - def _reciprocal(self, node: fx.Node) -> relax.Var: - x = self.env[node.args[0]] - return self.block_builder.emit(relax.op.divide(relax.const(1.0, x.ty.dtype.dtype), x)) - def _sqrt(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] dtype = x.ty.dtype.dtype @@ -1850,8 +1846,8 @@ def create_convert_map( "bitwise_xor.Scalar": self._binary_op(relax.op.bitwise_xor, operator.xor), "bitwise_or_.Tensor": self._binary_op(relax.op.bitwise_or, operator.or_), "bitwise_or.Tensor": self._binary_op(relax.op.bitwise_or, operator.or_), - "div.Scalar": self._binary_op(relax.op.divide, operator.truediv), - "div.Tensor": self._binary_op(relax.op.divide, operator.truediv), + "div.Scalar": self._true_divide, + "div.Tensor": self._true_divide, "div.Tensor_mode": self._div, "eq.Scalar": self._binary_op(relax.op.equal, operator.eq), "eq.Tensor": self._binary_op(relax.op.equal, operator.eq), diff --git a/python/tvm/relax/frontend/torch/fx_translator.py b/python/tvm/relax/frontend/torch/fx_translator.py index 2e35ce6ce704..9df3438c1d76 100644 --- a/python/tvm/relax/frontend/torch/fx_translator.py +++ b/python/tvm/relax/frontend/torch/fx_translator.py @@ -64,10 +64,6 @@ def _fetch_attr(self, model, target: str): ########## Unary Ops ########## - def _reciprocal(self, node: fx.Node) -> relax.Var: - x = self.env[node.args[0]] - return self.block_builder.emit(relax.op.divide(relax.const(1.0, x.ty.dtype), x)) - def _leakyrelu_module(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] module = self.named_modules[node.target] @@ -942,7 +938,7 @@ def create_convert_map( "rshift": self._binary_op(relax.op.right_shift, operator.rshift), "rsub": self._rsub, "sub": self._binary_op(relax.op.subtract, operator.sub), - "truediv": self._binary_op(relax.op.divide, operator.truediv), + "truediv": self._true_divide, "xor": self._binary_op(relax.op.bitwise_xor, operator.xor), # neural network "adaptive_avg_pool1d": self._adaptive_avg_pool1d, diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index d78629f5737b..ad0b356525c7 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -1435,6 +1435,177 @@ def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tuple( verify_model(BinaryPromoteRHS(), example_args, {}, expected_promote_rhs) +def test_binary_python_scalar_promotes_the_tensor(): + # A Python scalar only widens the tensor when its category is higher, and then the + # constant has to be built in the promoted dtype. Truncating the scalar to the + # tensor's dtype instead turns ``x + 0.5`` on an int64 tensor into ``x + 0``. + class AddHalf(Module): + def forward(self, x): + return x + 0.5 + + @tvm.script.ir_module + class expected_add_half: + @R.function + def main(x: R.Tensor((3,), dtype="int64")) -> R.Tuple(R.Tensor((3,), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((3,), dtype="float32") = R.astype(x, dtype="float32") + lv1: R.Tensor((3,), dtype="float32") = R.add(lv, R.const(0.5, "float32")) + gv: R.Tuple(R.Tensor((3,), dtype="float32")) = (lv1,) + R.output(gv) + return gv + + verify_model(AddHalf(), (torch.tensor([1, 2, 3]),), {}, expected_add_half) + + +def _scalar_promotion_cases(): + """(tensor dtype, scalar): torch.result_type decides the outcome in every case.""" + return [ + (torch.int64, 0.5), + (torch.int32, 1.5), + (torch.uint8, -0.5), + (torch.int64, 2), + (torch.float16, 2), + (torch.float32, True), + (torch.bool, 2), + (torch.bool, 0.5), + ] + + +def _scalar_input(dtype): + if dtype is torch.bool: + return torch.tensor([True, False, True]) + return torch.tensor([1, 2, 3], dtype=dtype) + + +def _verify_scalar_promotion(model, x): + # Compare both the result dtype and the values: the failure mode this guards is + # ``int * 0.5`` coming back as an int64 tensor of zeros, which a shape check passes. + with torch.no_grad(): + want = model(x) + mod = from_exported_program(export(model, (x,))) + got_dtype = str(mod["main"].ret_ty.fields[0].dtype) + assert got_dtype == str(want.dtype).replace("torch.", ""), ( + f"result dtype {got_dtype}, torch gives {want.dtype}" + ) + verify_model_numerically(model, (x,), rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("op", [operator.add, operator.mul, operator.lt, operator.ge, operator.eq]) +@pytest.mark.parametrize("dtype, scalar", _scalar_promotion_cases()) +def test_binary_python_scalar_promotion_values(op, dtype, scalar): + class Scalar(Module): + def forward(self, x): + return op(x, scalar) + + class ScalarOnTheLeft(Module): + def forward(self, x): + return op(scalar, x) + + x = _scalar_input(dtype) + _verify_scalar_promotion(Scalar(), x) + _verify_scalar_promotion(ScalarOnTheLeft(), x) + + +@pytest.mark.parametrize( + "dtype, scalar", [(torch.int64, 0.5), (torch.int32, 1.5), (torch.int64, 2)] +) +def test_binary_python_scalar_promotion_sub_pow_remainder(dtype, scalar): + # These three reject a bool tensor in torch, so they get their own case list. + class Sub(Module): + def forward(self, x): + return x - scalar + + class RSub(Module): + def forward(self, x): + return scalar - x + + class Pow(Module): + def forward(self, x): + return x**scalar + + class Remainder(Module): + def forward(self, x): + return x % scalar + + x = _scalar_input(dtype) + for model in (Sub(), RSub(), Pow(), Remainder()): + _verify_scalar_promotion(model, x) + + +def test_true_division_of_integers_gives_float(): + # torch's `/` always produces a floating result: int64 / 2 is float32, not a + # truncating integer quotient. That rule sits on top of scalar promotion (an int + # scalar alone would not widen an int tensor), so it has its own converter. + class Div(Module): + def forward(self, x): + return x / 2 + + @tvm.script.ir_module + class expected_div: + @R.function + def main(x: R.Tensor((3,), dtype="int64")) -> R.Tuple(R.Tensor((3,), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((3,), dtype="float32") = R.astype(x, dtype="float32") + lv1: R.Tensor((), dtype="float32") = R.astype(R.const(2, "int64"), dtype="float32") + lv2: R.Tensor((3,), dtype="float32") = R.divide(lv, lv1) + gv: R.Tuple(R.Tensor((3,), dtype="float32")) = (lv2,) + R.output(gv) + return gv + + verify_model(Div(), (torch.tensor([3, 4, 5]),), {}, expected_div) + + +@pytest.mark.parametrize( + "dtype, scalar", + [(torch.int64, 2), (torch.int32, 3), (torch.uint8, 2), (torch.bool, 2), (torch.int64, 2.5)], +) +def test_true_division_values(dtype, scalar): + class Div(Module): + def forward(self, x): + return x / scalar + + class RDiv(Module): + def forward(self, x): + return scalar / x + + class DivTensor(Module): + def forward(self, x): + return x / (x + 1) + + x = ( + torch.tensor([3, 4, 5], dtype=dtype) + if dtype is not torch.bool + else torch.tensor([True, True, False]) + ) + for model in (Div(), DivTensor()) + ((RDiv(),) if dtype is not torch.bool else ()): + _verify_scalar_promotion(model, x) + + +@pytest.mark.parametrize("dtype", [torch.int64, torch.int32, torch.float32]) +def test_division_with_rounding_mode(dtype): + # `//` and torch.div(..., rounding_mode=...) keep the promoted dtype -- an integer + # pair stays integer -- and the negative inputs tell floor and trunc apart. + class FloorDiv(Module): + def forward(self, x): + return x // 2 + + class DivFloor(Module): + def forward(self, x): + return torch.div(x, 2, rounding_mode="floor") + + class DivTrunc(Module): + def forward(self, x): + return torch.div(x, 2, rounding_mode="trunc") + + class DivFloorTensor(Module): + def forward(self, x): + return torch.div(x, x - 4, rounding_mode="floor") + + x = torch.tensor([-7, -3, 3, 7], dtype=dtype) + for model in (FloorDiv(), DivFloor(), DivTrunc(), DivFloorTensor()): + _verify_scalar_promotion(model, x) + + operator_binary_2 = [ (operator.eq, R.equal), (operator.ne, R.not_equal), @@ -8120,16 +8291,20 @@ def main(input: R.Tensor((9, 9), dtype="float32")) -> R.Tuple( lv: R.Tensor((9,), dtype="int64") = R.arange( R.prim_value(0), R.prim_value(9), R.prim_value(1), dtype="int64" ) - lv1: R.Tensor((9,), dtype="bool") = R.less(lv, R.const(4, "int64")) - lv2: R.Tensor((9,), dtype="float32") = R.astype(lv, dtype="float32") - lv3: R.Tensor((9,), dtype="float32") = R.multiply(lv2, R.const(0.125, "float32")) - lv4: R.Tensor((9,), dtype="float32") = R.add(lv3, R.const(0.0, "float32")) - lv5: R.Tensor((9,), dtype="int64") = R.subtract(R.const(8, "int64"), lv) - lv6: R.Tensor((9,), dtype="float32") = R.astype(lv5, dtype="float32") - lv7: R.Tensor((9,), dtype="float32") = R.multiply(lv6, R.const(0.125, "float32")) - lv8: R.Tensor((9,), dtype="float32") = R.subtract(R.const(1.0, "float32"), lv7) - lv9: R.Tensor((9,), dtype="float32") = R.where(lv1, lv4, lv8) - gv: R.Tuple(R.Tensor((9,), dtype="float32")) = (lv9,) + # torch's decomposition splits the range at ``i < 4.5``: the float + # scalar promotes the int64 index to float32 rather than being + # truncated to ``i < 4``. + lv1: R.Tensor((9,), dtype="float32") = R.astype(lv, dtype="float32") + lv2: R.Tensor((9,), dtype="bool") = R.less(lv1, R.const(4.5, "float32")) + lv3: R.Tensor((9,), dtype="float32") = R.astype(lv, dtype="float32") + lv4: R.Tensor((9,), dtype="float32") = R.multiply(lv3, R.const(0.125, "float32")) + lv5: R.Tensor((9,), dtype="float32") = R.add(lv4, R.const(0.0, "float32")) + lv6: R.Tensor((9,), dtype="int64") = R.subtract(R.const(8, "int64"), lv) + lv7: R.Tensor((9,), dtype="float32") = R.astype(lv6, dtype="float32") + lv8: R.Tensor((9,), dtype="float32") = R.multiply(lv7, R.const(0.125, "float32")) + lv9: R.Tensor((9,), dtype="float32") = R.subtract(R.const(1.0, "float32"), lv8) + lv10: R.Tensor((9,), dtype="float32") = R.where(lv2, lv5, lv9) + gv: R.Tuple(R.Tensor((9,), dtype="float32")) = (lv10,) R.output(gv) return gv