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 85a0d956ca1d..0bf1b7f58169 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -572,23 +572,69 @@ def call_binary_op(op, lhs, rhs): def _pow(self, node: fx.Node) -> relax.Var: lhs, rhs = self.retrieve_args(node) - # torch integer pow returns an integer tensor, but relax.op.power legalizes to - # TOPI power which requires floating-point inputs. Decompose an integer base with - # a constant non-negative integer exponent into repeated multiplication instead. - if ( - isinstance(lhs, relax.Expr) - and isinstance(lhs.ty, relax.TensorType) - and lhs.ty.dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT) - and isinstance(rhs, int) - and not isinstance(rhs, bool) - and rhs >= 0 - ): - if rhs == 0: - return self.block_builder.emit(relax.op.ones_like(lhs)) - result = lhs - for _ in range(rhs - 1): - result = self.block_builder.emit(relax.op.multiply(result, lhs)) - return result + if isinstance(lhs, relax.Expr) and isinstance(lhs.ty, relax.TensorType): + lhs_dtype = lhs.ty.dtype + is_integer_base = lhs_dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT) + is_float_base = lhs_dtype.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT) + + # A Python float promotes an integer tensor to PyTorch's default floating-point + # dtype. ExportedProgram records the inferred dtype, while plain FX does not. + if is_integer_base and isinstance(rhs, float): + output_meta = node.meta.get("val") + output_dtype = self._convert_data_type( + output_meta.dtype + if isinstance(output_meta, self.torch.Tensor) + else self.torch.get_default_dtype() + ) + lhs = self.block_builder.emit(relax.op.astype(lhs, output_dtype)) + lhs_dtype = lhs.ty.dtype + is_integer_base = False + is_float_base = True + + # Match the scalar conversion used by PyTorch's floating-point power kernels. + exponent_dtype = { + "float16": self.torch.float16, + "bfloat16": self.torch.bfloat16, + "float32": self.torch.float64, + "float64": self.torch.float64, + }.get(str(lhs_dtype)) + if ( + is_float_base + and exponent_dtype is not None + and isinstance(rhs, int | float) + and not isinstance(rhs, bool) + ): + rhs = self.torch.scalar_tensor(rhs, dtype=exponent_dtype, device="cpu").item() + + is_nonnegative_integral_exponent = ( + isinstance(rhs, int) and not isinstance(rhs, bool) and rhs >= 0 + ) or (isinstance(rhs, float) and rhs >= 0 and rhs.is_integer()) + + # TOPI power requires floating-point inputs, and some backends do not preserve + # the sign of a negative base for integral exponents. Decompose after applying + # PyTorch's dtype promotion so Python integer and float scalars behave alike. + if (is_integer_base or is_float_base) and is_nonnegative_integral_exponent: + exponent = int(rhs) + if exponent == 0: + return self.block_builder.emit(relax.op.ones_like(lhs)) + + # Exponentiation by squaring avoids linear graph growth for large exponents. + result = None + factor = lhs + while exponent: + if exponent & 1: + result = ( + factor + if result is None + else self.block_builder.emit(relax.op.multiply(result, factor)) + ) + exponent >>= 1 + if exponent: + factor = self.block_builder.emit(relax.op.multiply(factor, factor)) + return result + + if is_float_base and isinstance(rhs, float): + return self.block_builder.emit(relax.op.power(lhs, relax.const(rhs, lhs_dtype))) return self._binary_op(relax.op.power, operator.pow)(node) def _div(self, node: fx.Node) -> relax.Var: diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index 3799b8ed95f5..2d443aea1f7c 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -1079,9 +1079,8 @@ def main(input: R.Tensor((4,), dtype="int64")) -> R.Tuple(R.Tensor((4,), dtype=" # block 0 with R.dataflow(): lv: R.Tensor((4,), dtype="int64") = R.multiply(input, input) - lv1: R.Tensor((4,), dtype="int64") = R.multiply(lv, input) - lv2: R.Tensor((4,), dtype="int64") = R.multiply(lv1, input) - gv: R.Tuple(R.Tensor((4,), dtype="int64")) = (lv2,) + lv1: R.Tensor((4,), dtype="int64") = R.multiply(lv, lv) + gv: R.Tuple(R.Tensor((4,), dtype="int64")) = (lv1,) R.output(gv) return gv @@ -1089,6 +1088,132 @@ def main(input: R.Tensor((4,), dtype="int64")) -> R.Tuple(R.Tensor((4,), dtype=" verify_model(Pow(), example_args, {}, expected) +@pytest.mark.parametrize("exponent", [3, 3.0]) +def test_pow_float_integer_exponent(exponent): + class Pow(Module): + def forward(self, input): + return input.pow(exponent) + + @tvm.script.ir_module + class expected: + @R.function + def main( + input: R.Tensor((4,), dtype="float32"), + ) -> R.Tuple(R.Tensor((4,), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((4,), dtype="float32") = R.multiply(input, input) + lv1: R.Tensor((4,), dtype="float32") = R.multiply(input, lv) + gv: R.Tuple(R.Tensor((4,), dtype="float32")) = (lv1,) + R.output(gv) + return gv + + example_args = (torch.tensor([-2.0, -1.0, 1.0, 2.0], dtype=torch.float32),) + verify_model(Pow(), example_args, {}, expected) + verify_model_numerically(Pow(), example_args) + + +@pytest.mark.parametrize("exponent", [0, 0.0, 1, 1.0]) +def test_pow_float_integer_exponent_identity_cases(exponent): + class Pow(Module): + def forward(self, input): + return input.pow(exponent) + + example_args = (torch.tensor([-2.0, -1.0, 1.0, 2.0], dtype=torch.float32),) + verify_model_numerically(Pow(), example_args) + + +def test_pow_float_integer_exponent_large(): + class Pow(Module): + def forward(self, input): + return input.pow(17.0) + + @tvm.script.ir_module + class expected: + @R.function + def main(input: R.Tensor((2,), dtype="float32")) -> R.Tuple( + R.Tensor((2,), dtype="float32") + ): + with R.dataflow(): + lv: R.Tensor((2,), dtype="float32") = R.multiply(input, input) + lv1: R.Tensor((2,), dtype="float32") = R.multiply(lv, lv) + lv2: R.Tensor((2,), dtype="float32") = R.multiply(lv1, lv1) + lv3: R.Tensor((2,), dtype="float32") = R.multiply(lv2, lv2) + lv4: R.Tensor((2,), dtype="float32") = R.multiply(input, lv3) + gv: R.Tuple(R.Tensor((2,), dtype="float32")) = (lv4,) + R.output(gv) + return gv + + example_args = (torch.tensor([-1.25, 0.5], dtype=torch.float32),) + verify_model(Pow(), example_args, {}, expected) + verify_model_numerically(Pow(), example_args, rtol=1e-6, atol=1e-6) + + +def test_pow_integer_base_float_exponent(): + class Pow(Module): + def forward(self, input): + return input.pow(3.0) + + @tvm.script.ir_module + class expected: + @R.function + def main(input: R.Tensor((4,), dtype="int32")) -> R.Tuple(R.Tensor((4,), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((4,), dtype="float32") = R.astype(input, dtype="float32") + lv1: R.Tensor((4,), dtype="float32") = R.multiply(lv, lv) + lv2: R.Tensor((4,), dtype="float32") = R.multiply(lv, lv1) + gv: R.Tuple(R.Tensor((4,), dtype="float32")) = (lv2,) + R.output(gv) + return gv + + example_args = (torch.tensor([-2, -1, 1, 2], dtype=torch.int32),) + verify_model(Pow(), example_args, {}, expected) + verify_model_numerically(Pow(), example_args) + + +def test_pow_integer_base_fractional_exponent(): + class Pow(Module): + def forward(self, input): + return input.pow(0.5) + + @tvm.script.ir_module + class expected: + @R.function + def main(input: R.Tensor((4,), dtype="int32")) -> R.Tuple(R.Tensor((4,), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((4,), dtype="float32") = R.astype(input, dtype="float32") + lv1: R.Tensor((4,), dtype="float32") = R.power(lv, R.const(0.5, "float32")) + gv: R.Tuple(R.Tensor((4,), dtype="float32")) = (lv1,) + R.output(gv) + return gv + + example_args = (torch.tensor([1, 4, 9, 16], dtype=torch.int32),) + verify_model(Pow(), example_args, {}, expected) + verify_model_numerically(Pow(), example_args) + + +@pytest.mark.parametrize( + "dtype, exponent", + [ + (torch.float16, 2049.0), + (torch.bfloat16, 257.0), + (torch.float32, 2**53 + 1), + (torch.float64, 2**53 + 1), + ], +) +def test_pow_float_exponent_rounding(dtype, exponent): + class Pow(Module): + def forward(self, input): + return input.pow(exponent) + + input = torch.tensor([-1.0], dtype=dtype) + expected = Pow()(input) + mod = from_exported_program(export(Pow(), args=(input,))) + vm = relax.VirtualMachine(relax.build(mod, "llvm"), tvm.cpu()) + actual = torch.from_dlpack(vm["main"](tvm.runtime.from_dlpack(input))[0]) + + torch.testing.assert_close(actual, expected) + + def test_logsoftmax(): class LogSoftmax(Module): def forward(self, input): @@ -1336,17 +1461,36 @@ def __init__(self, op): def forward(self, lhs): return self.op(lhs, 1.0) - @tvm.script.ir_module - class expected_binary2: - @R.function - def main( - lhs: R.Tensor((10, 10), dtype="float32"), - ) -> R.Tuple(R.Tensor((10, 10), dtype="float32")): - with R.dataflow(): - lv: R.Tensor((10, 10), dtype="float32") = relax_op(lhs, R.const(1.0)) - gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv,) - R.output(gv) - return gv + if op is operator.pow: + + @tvm.script.ir_module + class expected_power: + @R.function + def main( + lhs: R.Tensor((10, 10), dtype="float32"), + ) -> R.Tuple(R.Tensor((10, 10), dtype="float32")): + with R.dataflow(): + gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lhs,) + R.output(gv) + return gv + + expected_binary2 = expected_power + + else: + + @tvm.script.ir_module + class expected_other_binary: + @R.function + def main( + lhs: R.Tensor((10, 10), dtype="float32"), + ) -> R.Tuple(R.Tensor((10, 10), dtype="float32")): + with R.dataflow(): + lv: R.Tensor((10, 10), dtype="float32") = relax_op(lhs, R.const(1.0)) + gv: R.Tuple(R.Tensor((10, 10), dtype="float32")) = (lv,) + R.output(gv) + return gv + + expected_binary2 = expected_other_binary # In-place ops (add_, mul_, ...) produce the same Relax program as their # functional counterparts: mutation outputs are dropped by the importer. diff --git a/tests/python/relax/test_frontend_from_fx.py b/tests/python/relax/test_frontend_from_fx.py index a489977958c7..e4eae1d17cda 100644 --- a/tests/python/relax/test_frontend_from_fx.py +++ b/tests/python/relax/test_frontend_from_fx.py @@ -2241,17 +2241,36 @@ def __init__(self, op): def forward(self, lhs): return self.op(lhs, 1.0) - @tvm.script.ir_module - class expected_binary2: - @R.function - def main( - lhs: R.Tensor((1, 3, 10, 10), dtype="float32"), - ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): - with R.dataflow(): - lv: R.Tensor((1, 3, 10, 10), dtype="float32") = relax_op(lhs, R.const(1.0)) - gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv - R.output(gv) - return gv + if op is operator.pow: + + @tvm.script.ir_module + class expected_power: + @R.function + def main( + lhs: R.Tensor((1, 3, 10, 10), dtype="float32"), + ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): + with R.dataflow(): + gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lhs + R.output(gv) + return gv + + expected_binary2 = expected_power + + else: + + @tvm.script.ir_module + class expected_other_binary: + @R.function + def main( + lhs: R.Tensor((1, 3, 10, 10), dtype="float32"), + ) -> R.Tensor((1, 3, 10, 10), dtype="float32"): + with R.dataflow(): + lv: R.Tensor((1, 3, 10, 10), dtype="float32") = relax_op(lhs, R.const(1.0)) + gv: R.Tensor((1, 3, 10, 10), dtype="float32") = lv + R.output(gv) + return gv + + expected_binary2 = expected_other_binary verify_model(Binary1(op), input_info1, {}, expected_binary1) verify_model(Binary2(op), input_info2, {}, expected_binary2) @@ -3578,9 +3597,73 @@ class expected: def main(inp_0: R.Tensor((4,), dtype="int64")) -> R.Tensor((4,), dtype="int64"): with R.dataflow(): lv: R.Tensor((4,), dtype="int64") = R.multiply(inp_0, inp_0) - lv1: R.Tensor((4,), dtype="int64") = R.multiply(lv, inp_0) - lv2: R.Tensor((4,), dtype="int64") = R.multiply(lv1, inp_0) - gv: R.Tensor((4,), dtype="int64") = lv2 + lv1: R.Tensor((4,), dtype="int64") = R.multiply(lv, lv) + gv: R.Tensor((4,), dtype="int64") = lv1 + R.output(gv) + return gv + + verify_model(Pow(), input_info, {}, expected) + + +@pytest.mark.parametrize("exponent", [3, 3.0]) +def test_pow_float_integer_exponent(exponent): + input_info = [([4], "float32")] + + class Pow(Module): + def forward(self, input): + return input.pow(exponent) + + @tvm.script.ir_module + class expected: + @R.function + def main(inp_0: R.Tensor((4,), dtype="float32")) -> R.Tensor((4,), dtype="float32"): + with R.dataflow(): + lv: R.Tensor((4,), dtype="float32") = R.multiply(inp_0, inp_0) + lv1: R.Tensor((4,), dtype="float32") = R.multiply(inp_0, lv) + gv: R.Tensor((4,), dtype="float32") = lv1 + R.output(gv) + return gv + + verify_model(Pow(), input_info, {}, expected) + + +def test_pow_integer_base_float_exponent(): + input_info = [([4], "int32")] + + class Pow(Module): + def forward(self, input): + return input.pow(3.0) + + @tvm.script.ir_module + class expected: + @R.function + def main(inp_0: R.Tensor((4,), dtype="int32")) -> R.Tensor((4,), dtype="float32"): + with R.dataflow(): + lv: R.Tensor((4,), dtype="float32") = R.astype(inp_0, dtype="float32") + lv1: R.Tensor((4,), dtype="float32") = R.multiply(lv, lv) + lv2: R.Tensor((4,), dtype="float32") = R.multiply(lv, lv1) + gv: R.Tensor((4,), dtype="float32") = lv2 + R.output(gv) + return gv + + verify_model(Pow(), input_info, {}, expected) + + +def test_pow_integer_base_fractional_exponent(): + input_info = [([4], "int32")] + + class Pow(Module): + def forward(self, input): + return input.pow(0.5) + + @tvm.script.ir_module + class expected: + @R.function + def main(inp_0: R.Tensor((4,), dtype="int32")) -> R.Tensor((4,), dtype="float32"): + with R.dataflow(): + lv: R.Tensor((4,), dtype="float32") = R.astype(inp_0, dtype="float32") + lv1: R.Tensor((4,), dtype="float32") = R.power(lv, R.const(0.5, "float32")) + gv: R.Tensor((4,), dtype="float32") = lv1 R.output(gv) return gv