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 b0bb14ac95ad..d600987cdd7b 100644 --- a/python/tvm/relax/frontend/torch/base_fx_graph_translator.py +++ b/python/tvm/relax/frontend/torch/base_fx_graph_translator.py @@ -2501,15 +2501,35 @@ def _eye(self, node: fx.Node) -> relax.Var: def _fill(self, node: fx.Node) -> relax.Var: args = self.retrieve_args(node) x = args[0] - dtype = x.ty.dtype - value = args[1] if isinstance(args[1], relax.Expr) else relax.const(args[1], dtype) + dtype = str(x.ty.dtype) + value = self._convert_scalar_fill_value(args[1], dtype) return self.block_builder.emit(relax.op.full(x.ty.shape, value, dtype)) + def _prim_value_to_scalar_tensor(self, value: relax.Expr, dtype: str) -> relax.Var: + """Materialize an integer or boolean primitive value as a rank-zero tensor.""" + if not value.ty.matches_code(DataTypeCode.INT, DataTypeCode.UINT, DataTypeCode.BOOL): + raise TypeError(f"Cannot materialize primitive value of dtype {value.ty} as a tensor") + shape_value = value if str(value.ty) == "int64" else value.astype("int64") + value_tensor = self.block_builder.emit( + relax.op.shape_to_tensor(relax.ShapeExpr([shape_value])) + ) + if dtype != "int64": + value_tensor = self.block_builder.emit(relax.op.astype(value_tensor, dtype)) + return self.block_builder.emit(relax.op.squeeze(value_tensor, axis=[0])) + + def _convert_scalar_fill_value(self, value, dtype: str) -> relax.Expr: + """Convert a PyTorch scalar fill value to a rank-zero Relax tensor.""" + if isinstance(getattr(value, "ty", None), PrimType): + return self._prim_value_to_scalar_tensor(value, dtype) + if isinstance(value, relax.Expr): + return value + return relax.const(value, dtype) + def _inplace_fill(self, node: fx.Node) -> relax.Var: args = self.retrieve_args(node) x = args[0] - dtype = x.ty.dtype.dtype - value = args[1] if isinstance(args[1], relax.Expr) else relax.const(args[1], dtype) + dtype = str(x.ty.dtype) + value = self._convert_scalar_fill_value(args[1], dtype) filled = self.block_builder.emit(relax.op.full(x.ty.shape, value, dtype)) self.env[node.args[0]] = filled return filled @@ -2519,10 +2539,23 @@ def _full(self, node: fx.Node) -> relax.Var: args = self.retrieve_args(node) size = relax.ShapeExpr(args[0] if isinstance(args[0], list | tuple) else (args[0],)) - dtype = self._convert_data_type( - node.kwargs.get("dtype", torch.get_default_dtype()), self.env - ) - value = args[1] if isinstance(args[1], relax.expr.Constant) else relax.const(args[1], dtype) + torch_dtype = node.kwargs.get("dtype") + if torch_dtype is None: + output_meta = node.meta.get("val") + if output_meta is None: + output_meta = node.meta.get("tensor_meta") + torch_dtype = getattr(output_meta, "dtype", None) + if torch_dtype is None: + if isinstance(args[1], bool): + torch_dtype = "bool" + elif isinstance(args[1], int): + torch_dtype = "int64" + elif isinstance(getattr(args[1], "ty", None), PrimType): + torch_dtype = args[1].ty.dtype + else: + torch_dtype = torch.get_default_dtype() + dtype = self._convert_data_type(torch_dtype, self.env) + value = self._convert_scalar_fill_value(args[1], dtype) return self.block_builder.emit( relax.op.full( size, @@ -2532,15 +2565,28 @@ def _full(self, node: fx.Node) -> relax.Var: ) def _full_like(self, node: fx.Node) -> relax.Var: - x = self.env[node.args[0]] - value = node.args[1] - fill_value = relax.const(value) - - x_dtype = x.ty.dtype.dtype - fill_dtype = None - if isinstance(value, int | float) and (math.isinf(value) or math.isnan(value)): + args = self.retrieve_args(node) + x = args[0] + value = args[1] + x_dtype = str(x.ty.dtype) + torch_dtype = node.kwargs.get("dtype") + dtype = self._convert_data_type(x_dtype if torch_dtype is None else torch_dtype, self.env) + fill_dtype = dtype if dtype != x_dtype else None + if ( + fill_dtype is None + and isinstance(value, int | float) + and (math.isinf(value) or math.isnan(value)) + ): if not ("float" in x_dtype or "bfloat16" in x_dtype): fill_dtype = "float32" + dtype = fill_dtype + + if isinstance(getattr(value, "ty", None), PrimType): + fill_value = self._prim_value_to_scalar_tensor(value, dtype) + elif isinstance(value, relax.Expr): + fill_value = value + else: + fill_value = relax.const(value, dtype) return self.block_builder.emit(relax.op.full_like(x, fill_value, dtype=fill_dtype)) @@ -2551,16 +2597,14 @@ def _index_select(self, node: fx.Node) -> relax.Var: return self.block_builder.emit(relax.op.take(x, index, dim)) def _inplace_masked_fill(self, node: fx.Node) -> relax.Var: - x = self.env[node.args[0]] - mask = self.env[node.args[1]] - value = node.args[2] - rx_value = relax.const(value) - - x_dtype = x.ty.dtype.dtype + args = self.retrieve_args(node) + x, mask, value = args[:3] + x_dtype = str(x.ty.dtype) fill_dtype = None if isinstance(value, int | float) and (math.isinf(value) or math.isnan(value)): if not ("float" in x_dtype or "bfloat16" in x_dtype): fill_dtype = "float32" + rx_value = self._convert_scalar_fill_value(value, fill_dtype or x_dtype) values = self.block_builder.emit(relax.op.full_like(x, rx_value, dtype=fill_dtype)) @@ -2596,16 +2640,14 @@ def _linspace(self, node: fx.Node) -> relax.Var: ) def _masked_fill(self, node: fx.Node) -> relax.Var: - x = self.env[node.args[0]] - mask = self.env[node.args[1]] - value = node.args[2] - rx_value = relax.const(value) - - x_dtype = x.ty.dtype.dtype + args = self.retrieve_args(node) + x, mask, value = args[:3] + x_dtype = str(x.ty.dtype) fill_dtype = None if isinstance(value, int | float) and (math.isinf(value) or math.isnan(value)): if not ("float" in x_dtype or "bfloat16" in x_dtype): fill_dtype = "float32" + rx_value = self._convert_scalar_fill_value(value, fill_dtype or x_dtype) values = self.block_builder.emit(relax.op.full_like(x, rx_value, dtype=fill_dtype)) @@ -2807,6 +2849,23 @@ def _getitem(self, node: fx.Node) -> relax.Var: def _item(self, node: fx.Node) -> relax.Var: x = self.env[node.args[0]] + shape = self.shape_of(x) + dtype = x.ty.dtype + analyzer = tvm.arith.Analyzer() + has_single_element = shape is not None and all( + analyzer.can_prove_equal(dim, 1) for dim in shape + ) + if has_single_element and dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT): + scalar = x + if str(dtype) != "int64": + scalar = self.block_builder.emit(relax.op.astype(scalar, "int64")) + scalar = self.block_builder.emit(relax.op.reshape(scalar, [1])) + shape_value = self.block_builder.emit(relax.op.tensor_to_shape(scalar)) + dim = tirx.Var(f"{node.name}_dim", "int64") + self.block_builder.match_cast(shape_value, relax.ShapeType([dim])) + return dim + if shape is not None and len(shape) == 0: + return x return self.block_builder.emit(relax.op.take(x, relax.const(0, "int64"), axis=0)) def _sym_size_int(self, node: fx.Node) -> relax.Expr: diff --git a/python/tvm/relax/frontend/torch/exported_program_translator.py b/python/tvm/relax/frontend/torch/exported_program_translator.py index 0e69074af601..ced0aa7b28bd 100644 --- a/python/tvm/relax/frontend/torch/exported_program_translator.py +++ b/python/tvm/relax/frontend/torch/exported_program_translator.py @@ -1481,8 +1481,12 @@ def _as_strided(self, node: fx.Node) -> relax.Var: ########## Symbolic Shape Constraints ########## - def _symbolic_comparison(self, _: fx.Node) -> relax.Expr: - return self.block_builder.emit(relax.const(True, dtype="bool")) + def _symbolic_comparison(self, intrinsic_op: Callable) -> Callable: + def convert(node: fx.Node) -> relax.Expr: + lhs, rhs = self.retrieve_args(node) + return self.block_builder.emit(relax.prim_value(intrinsic_op(lhs, rhs))) + + return convert ########## Higher-Order Ops ########## @@ -1754,6 +1758,7 @@ def create_convert_map( "logical_xor.default": self._logical_xor, "log_softmax.int": self._log_softmax, "_log_softmax.default": self._log_softmax, + "neg": lambda node: operator.neg(self.retrieve_args(node)[0]), "neg.default": self._unary_op(relax.op.negative), "pad.default": self._pad, "constant_pad_nd.default": self._constant_pad_nd, @@ -1792,6 +1797,7 @@ def create_convert_map( "triu.default": self._tril_triu(relax.op.triu), "trunc.default": self._unary_op(relax.op.trunc), # binary + "add": self._binary_op(relax.op.add, operator.add), "add.Tensor": self._binary_op(relax.op.add, operator.add), "add.Scalar": self._binary_op(relax.op.add, operator.add), "add_.Tensor": self._binary_op(relax.op.add, operator.add), @@ -1810,6 +1816,7 @@ def create_convert_map( "eq.Scalar": self._binary_op(relax.op.equal, operator.eq), "eq.Tensor": self._binary_op(relax.op.equal, operator.eq), "floor_divide.default": self._binary_op(relax.op.floor_divide, operator.floordiv), + "floordiv": self._binary_op(relax.op.floor_divide, operator.floordiv), "fmod.Scalar": self._fmod, "fmod.Tensor": self._fmod, "logaddexp.default": self._binary_op(relax.op.log_add_exp, torch.logaddexp), @@ -1835,6 +1842,7 @@ def create_convert_map( "minimum.default": self._binary_op(relax.op.minimum, torch.minimum), "remainder.Tensor": self._binary_op(relax.op.floor_mod, operator.mod), "remainder.Scalar": self._binary_op(relax.op.floor_mod, operator.mod), + "mod": self._binary_op(relax.op.floor_mod, operator.mod), "mul": self._binary_op(relax.op.multiply, operator.mul), "mul.Tensor": self._binary_op(relax.op.multiply, operator.mul), "mul.Scalar": self._binary_op(relax.op.multiply, operator.mul), @@ -1847,6 +1855,7 @@ def create_convert_map( "pow.Scalar": self._binary_op(relax.op.power, operator.pow), "pow.Tensor_Scalar": self._pow, "pow.Tensor_Tensor": self._binary_op(relax.op.power, operator.pow), + "sub": self._binary_op(relax.op.subtract, operator.sub), "sub.Tensor": self._binary_op(relax.op.subtract, operator.sub), "sub.Scalar": self._binary_op(relax.op.subtract, operator.sub), "__and__.Tensor": self._binary_op(relax.op.bitwise_and, operator.and_), @@ -2031,13 +2040,15 @@ def create_convert_map( "item.default": self._item, "sym_size.int": self._sym_size_int, "_local_scalar_dense.default": self._item, - # symbolic shape constraints (no-ops for compilation) + # symbolic shape operations and constraints "sym_constrain_range_for_size.default": lambda node: self.env[node.args[0]], "_assert_scalar.default": lambda node: self.env[node.args[0]], - "ge": self._symbolic_comparison, - "le": self._symbolic_comparison, - "gt": self._symbolic_comparison, - "lt": self._symbolic_comparison, + "ge": self._symbolic_comparison(operator.ge), + "le": self._symbolic_comparison(operator.le), + "gt": self._symbolic_comparison(operator.gt), + "lt": self._symbolic_comparison(operator.lt), + "eq": self._symbolic_comparison(operator.eq), + "ne": self._symbolic_comparison(operator.ne), # higher-order ops "cond": self._cond, } diff --git a/tests/python/relax/test_frontend_dynamo.py b/tests/python/relax/test_frontend_dynamo.py index 72e05a56494c..e795b0881f4b 100644 --- a/tests/python/relax/test_frontend_dynamo.py +++ b/tests/python/relax/test_frontend_dynamo.py @@ -504,7 +504,7 @@ def main( ) -> R.Tensor((256, 256), dtype="float32"): with R.dataflow(): lv: R.Tensor((256, 256), dtype="float32") = R.full_like( - inp_1, R.const(0, "int32"), dtype=None + inp_1, R.const(0.0, "float32"), dtype=None ) lv1: R.Tensor((256, 256), dtype="float32") = R.where(inp_0, lv, inp_1) gv: R.Tensor((256, 256), dtype="float32") = lv1 diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index e9d2ac8b704f..7dc3c7356414 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -5549,6 +5549,265 @@ def forward(self, x): assert tvm.arith.Analyzer().can_prove_equal(output_shape[2], input_shape[1]) +def test_dynamic_scalar_item_in_shape_operations(): + class DynamicShapeOps(torch.nn.Module): + def forward(self, x): + lengths = torch.full( + (x.shape[0],), + x.shape[1], + device=x.device, + dtype=torch.int64, + ) + max_len = lengths.max().item() + positions = torch.arange(max_len, device=x.device) + mask = positions.unsqueeze(0).expand(x.shape[0], -1) == lengths.unsqueeze(1) + filled = torch.full( + (x.shape[0], x.shape[1]), + x.shape[1], + device=x.device, + dtype=torch.int64, + ) + shifted = torch.arange(x.shape[1] + 1, device=x.device) + shortened = torch.full_like(lengths, x.shape[1] - 1) + return mask, filled, shifted, shortened + + example_args = (torch.randn(1, 4, 3, dtype=torch.float32),) + tokens = torch.export.Dim("tokens", min=1, max=8) + exported_program = export( + DynamicShapeOps(), + args=example_args, + dynamic_shapes={"x": {1: tokens}}, + ) + mod = from_exported_program(exported_program) + + script = mod.script() + assert "R.tensor_to_shape" in script + assert "R.shape_to_tensor" in script + + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + for token_count in (4, 6): + torch_input = torch.randn(1, token_count, 3) + expected = DynamicShapeOps()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + + +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +@pytest.mark.parametrize( + ("item_dtype", "item_shape"), + [ + (torch.int8, ()), + (torch.uint8, ()), + (torch.int16, ()), + (torch.int32, ()), + (torch.int64, ()), + (torch.int64, (1,)), + (torch.int64, (1, 1)), + ], +) +def test_dynamic_scalar_item_single_element_integer_dtypes(item_dtype, item_shape): + class DynamicItem(torch.nn.Module): + def forward(self, x): + lengths = torch.full( + (x.shape[0],), + x.shape[1], + device=x.device, + dtype=item_dtype, + ) + value = lengths.max().reshape(item_shape).item() + return torch.arange(value, device=x.device), torch.full_like( + x, value, dtype=torch.int64 + ) + + rows = torch.export.Dim("rows", min=1, max=8) + columns = torch.export.Dim("columns", min=1, max=8) + example_args = (torch.randn(3, 4, dtype=torch.float32),) + exported_program = export( + DynamicItem(), + args=example_args, + dynamic_shapes={"x": {0: rows, 1: columns}}, + ) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape in ((3, 4), (5, 2)): + torch_input = torch.randn(shape, dtype=torch.float32) + expected = DynamicItem()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + + +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_runtime_scalar_item_integer_value_ranges(): + class RuntimeItems(torch.nn.Module): + def forward(self, x, signed, unsigned, wide): + return ( + torch.full_like(x, signed.item(), dtype=torch.int64), + torch.full_like(x, unsigned.item(), dtype=torch.int64), + torch.full_like(x, wide.item(), dtype=torch.int64), + ) + + example_args = ( + torch.randn(2, 3, dtype=torch.float32), + torch.tensor([[-3]], dtype=torch.int8), + torch.tensor([200], dtype=torch.uint8), + torch.tensor(1 << 40, dtype=torch.int64), + ) + exported_program = export(RuntimeItems(), args=example_args) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + expected = RuntimeItems()(*example_args) + actual = vm["main"](*(tvm.runtime.tensor(arg.numpy()) for arg in example_args)) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + + +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_dynamic_scalar_fill_operations(): + class DynamicFills(torch.nn.Module): + def forward(self, x): + inferred_dtype = torch.full((x.shape[0],), x.shape[1], device=x.device) + explicit_dtype = torch.full_like(x, x.shape[0], dtype=torch.float64) + filled = torch.fill(x, x.shape[1]) + filled_inplace = x.clone() + filled_inplace.fill_(x.shape[0]) + return inferred_dtype, explicit_dtype, filled, filled_inplace + + example_args = (torch.randn(3, 4, dtype=torch.float32),) + rows = torch.export.Dim("rows", min=1, max=8) + columns = torch.export.Dim("columns", min=1, max=8) + exported_program = export( + DynamicFills(), + args=example_args, + dynamic_shapes={"x": {0: rows, 1: columns}}, + ) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape in ((3, 4), (5, 2)): + torch_input = torch.randn(shape, dtype=torch.float32) + expected = DynamicFills()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + actual_arrays = [value.numpy() for value in actual] + + assert actual_arrays[0].dtype == np.dtype("int64") + assert actual_arrays[1].dtype == np.dtype("float64") + for actual_value, expected_value in zip(actual_arrays, expected): + assert actual_value.dtype == expected_value.numpy().dtype + np.testing.assert_array_equal(actual_value, expected_value.numpy()) + + +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_dynamic_boolean_fill_operations(): + class DynamicBooleanFills(torch.nn.Module): + def forward(self, x): + value = x.shape[0] == x.shape[1] + base = x > 0 + filled = torch.fill(base, value) + filled_inplace = base.clone() + filled_inplace.fill_(value) + created = torch.full( + (x.shape[0], x.shape[1]), + value, + device=x.device, + dtype=torch.bool, + ) + like = torch.full_like(x, value, dtype=torch.bool) + return created, like, filled, filled_inplace + + rows = torch.export.Dim("rows", min=1, max=8) + columns = torch.export.Dim("columns", min=1, max=8) + example_args = (torch.randn(3, 4, dtype=torch.float32),) + exported_program = export( + DynamicBooleanFills(), + args=example_args, + dynamic_shapes={"x": {0: rows, 1: columns}}, + ) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape in ((3, 3), (3, 4)): + torch_input = torch.randn(shape, dtype=torch.float32) + expected = DynamicBooleanFills()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + + +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_dynamic_scalar_arithmetic(): + class DynamicScalarArithmetic(torch.nn.Module): + def forward(self, x): + lengths = torch.full( + (x.shape[0],), + x.shape[1], + device=x.device, + dtype=torch.int64, + ) + value = lengths.max().item() + return ( + torch.full_like(x, -value, dtype=torch.int64), + torch.arange(value // 2, device=x.device), + torch.arange(value % 3, device=x.device), + ) + + rows = torch.export.Dim("rows", min=1, max=8) + columns = torch.export.Dim("columns", min=1, max=8) + example_args = (torch.randn(3, 4, dtype=torch.float32),) + exported_program = export( + DynamicScalarArithmetic(), + args=example_args, + dynamic_shapes={"x": {0: rows, 1: columns}}, + ) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape in ((3, 4), (5, 2)): + torch_input = torch.randn(shape, dtype=torch.float32) + expected = DynamicScalarArithmetic()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + + +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +@pytest.mark.parametrize( + ("fill_value", "dtype"), + [ + (1 << 40, torch.int64), + (1.0 + 2**-40, torch.float64), + ], +) +def test_full_and_full_like_python_fill_value_uses_explicit_dtype(fill_value, dtype): + class FullLike(torch.nn.Module): + def forward(self, x): + return ( + torch.full(x.shape, fill_value, dtype=dtype), + torch.full_like(x, fill_value, dtype=dtype), + ) + + example_args = (torch.randn(2, 3, dtype=torch.float32),) + exported_program = export(FullLike(), args=example_args) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + expected = FullLike()(*example_args) + actual = vm["main"](tvm.runtime.tensor(example_args[0].numpy())) + for actual_value, expected_value in zip(actual, expected): + assert actual_value.numpy().dtype == expected_value.numpy().dtype + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + + def test_split(): class Chunk(Module): def forward(self, input): @@ -6072,6 +6331,71 @@ def main( verify_model(Masked_Fill_Inplace(), example_args, {}, Expected) +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_dynamic_scalar_masked_fill_operations(): + class DynamicMaskedFills(torch.nn.Module): + def forward(self, x): + mask = x > 0 + filled = x.masked_fill(mask, x.shape[0]) + filled_inplace = x.clone() + filled_inplace.masked_fill_(mask, x.shape[1]) + return filled, filled_inplace + + rows = torch.export.Dim("rows", min=1, max=8) + columns = torch.export.Dim("columns", min=1, max=8) + example_args = (torch.randn(3, 4, dtype=torch.float32),) + exported_program = export( + DynamicMaskedFills(), + args=example_args, + dynamic_shapes={"x": {0: rows, 1: columns}}, + ) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape in ((3, 4), (5, 2)): + torch_input = torch.randn(shape, dtype=torch.float32) + expected = DynamicMaskedFills()(torch_input) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy())) + for actual_value, expected_value in zip(actual, expected): + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + + +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +@pytest.mark.parametrize( + ("fill_value", "dtype"), + [ + (1 << 40, torch.int64), + (1.0 + 2**-40, torch.float64), + ], +) +def test_fill_and_masked_fill_python_value_uses_input_dtype(fill_value, dtype): + class MaskedFills(torch.nn.Module): + def forward(self, x, mask): + filled_all = torch.fill(x, fill_value) + filled_all_inplace = x.clone() + filled_all_inplace.fill_(fill_value) + filled = x.masked_fill(mask, fill_value) + filled_inplace = x.clone() + filled_inplace.masked_fill_(mask, fill_value) + return filled_all, filled_all_inplace, filled, filled_inplace + + example_args = ( + torch.arange(6, dtype=dtype).reshape(2, 3), + torch.tensor([[True, False, True], [False, True, False]]), + ) + exported_program = export(MaskedFills(), args=example_args) + mod = from_exported_program(exported_program, run_ep_decomposition=False) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + expected = MaskedFills()(*example_args) + actual = vm["main"](*(tvm.runtime.tensor(arg.numpy()) for arg in example_args)) + for actual_value, expected_value in zip(actual, expected): + assert actual_value.numpy().dtype == expected_value.numpy().dtype + np.testing.assert_array_equal(actual_value.numpy(), expected_value.numpy()) + + def test_masked_select(): class MaskedSelect(Module): def forward(self, data: torch.Tensor, mask: torch.Tensor): @@ -6094,8 +6418,8 @@ def main( ) lv4: R.Tensor((u0,), dtype="int64") = R.squeeze(lv3, axis=[0]) lv5: R.Tensor((u0,), dtype="float32") = R.take(lv, lv4, axis=0, mode="fast") - lv6: R.Tensor((), dtype="bool") = R.const(True, "bool") - lv7: R.Tensor((), dtype="bool") = R.const(True, "bool") + lv6: T.bool = u0 >= 0 + lv7: T.bool = u0 <= 6 gv: R.Tuple(R.Tensor((u0,), dtype="float32")) = (lv5,) R.output(gv) return gv @@ -6430,7 +6754,9 @@ def main(input: R.Tensor((128, 128), dtype="float32")) -> R.Tuple( R.Tensor((128, 128), dtype="float32") ): with R.dataflow(): - lv: R.Tensor((128, 128), dtype="float32") = R.full_like(input, R.const(1, "int32")) + lv: R.Tensor((128, 128), dtype="float32") = R.full_like( + input, R.const(1.0, "float32") + ) gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,) R.output(gv) return gv @@ -6452,7 +6778,9 @@ def main(input: R.Tensor((128, 128), dtype="float32")) -> R.Tuple( R.Tensor((128, 128), dtype="float32") ): with R.dataflow(): - lv: R.Tensor((128, 128), dtype="float32") = R.full_like(input, R.const(0, "int32")) + lv: R.Tensor((128, 128), dtype="float32") = R.full_like( + input, R.const(0.0, "float32") + ) gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,) R.output(gv) return gv @@ -6498,7 +6826,9 @@ def main(input: R.Tensor((128, 128), dtype="float32")) -> R.Tuple( R.Tensor((128, 128), dtype="float32") ): with R.dataflow(): - lv: R.Tensor((128, 128), dtype="float32") = R.full_like(input, R.const(0, "int32")) + lv: R.Tensor((128, 128), dtype="float32") = R.full_like( + input, R.const(0.0, "float32") + ) gv: R.Tuple(R.Tensor((128, 128), dtype="float32")) = (lv,) R.output(gv) return gv @@ -8692,7 +9022,7 @@ def main( s77 = T.int64() R.func_attr({"tir_var_lower_bound": {"s77": 1}}) cls = expected - gv: R.Tensor((), dtype="bool") = R.const(True, "bool") + gv: T.bool = s77 > 4 if gv: gv1: R.Tensor((s77, 4), dtype="float32") = cls.cond_true_branch_0(x) cond_result: R.Tensor((s77, 4), dtype="float32") = gv1 @@ -8712,6 +9042,64 @@ def main( ) +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_cond_shape_equality_predicate(): + class CondShapeEqualityModel(Module): + def forward(self, x): + def true_fn(x): + return x + 1.0 + + def false_fn(x): + return x - 1.0 + + return torch.cond(x.shape[0] == x.shape[1], true_fn, false_fn, (x,)) + + rows = torch.export.Dim("rows", min=1, max=8) + columns = torch.export.Dim("columns", min=1, max=8) + exported_program = export( + CondShapeEqualityModel(), + args=(torch.zeros(3, 3),), + dynamic_shapes={"x": {0: rows, 1: columns}}, + ) + mod = from_exported_program(exported_program) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape, expected_value in (((3, 3), 1.0), ((2, 3), -1.0)): + torch_input = torch.zeros(shape, dtype=torch.float32) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy()))[0] + np.testing.assert_array_equal(actual.numpy(), np.full(shape, expected_value, "float32")) + + +@pytest.mark.skipif(not env.has_llvm(), reason="need llvm") +def test_cond_shape_inequality_predicate(): + class CondShapeInequalityModel(Module): + def forward(self, x): + def true_fn(x): + return x + 1.0 + + def false_fn(x): + return x - 1.0 + + return torch.cond(x.shape[0] != x.shape[1], true_fn, false_fn, (x,)) + + rows = torch.export.Dim("rows", min=1, max=8) + columns = torch.export.Dim("columns", min=1, max=8) + exported_program = export( + CondShapeInequalityModel(), + args=(torch.zeros(2, 3),), + dynamic_shapes={"x": {0: rows, 1: columns}}, + ) + mod = from_exported_program(exported_program) + executable = relax.build(mod, tvm.target.Target("llvm")) + vm = relax.VirtualMachine(executable, tvm.cpu()) + + for shape, expected_value in (((2, 3), 1.0), ((3, 3), -1.0)): + torch_input = torch.zeros(shape, dtype=torch.float32) + actual = vm["main"](tvm.runtime.tensor(torch_input.numpy()))[0] + np.testing.assert_array_equal(actual.numpy(), np.full(shape, expected_value, "float32")) + + def test_cond_tuple_output(): """Cond where both branches return a tuple."""