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..7c16a4fe86d0 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.""" @@ -707,13 +739,25 @@ def promote_binary_op_args(lhs, rhs): return lhs, rhs elif isinstance(lhs, relax.Expr): assert isinstance(lhs.ty, relax.TensorType) - return lhs, relax.const(rhs, lhs.ty.dtype) + lhs, rhs = promote_scalar_operand(lhs, rhs) + return lhs, rhs elif isinstance(rhs, relax.Expr): assert isinstance(rhs.ty, relax.TensorType) - return relax.const(lhs, rhs.ty.dtype), rhs + rhs, lhs = promote_scalar_operand(rhs, lhs) + return lhs, rhs else: assert False + def promote_scalar_operand(tensor, scalar): + # 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 call_binary_op(op, lhs, rhs): lhs, rhs = promote_binary_op_args(lhs, rhs) return self.block_builder.emit(op(lhs, rhs)) @@ -725,9 +769,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 diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index d78629f5737b..28ef877fd9cb 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -1435,6 +1435,103 @@ 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) + + operator_binary_2 = [ (operator.eq, R.equal), (operator.ne, R.not_equal), @@ -8120,16 +8217,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