diff --git a/backends/arm/test/misc/test_int64.py b/backends/arm/test/misc/test_int64.py index 6bf6381b1ef..2ff584a8e01 100644 --- a/backends/arm/test/misc/test_int64.py +++ b/backends/arm/test/misc/test_int64.py @@ -139,16 +139,19 @@ def forward(self, cond, x): # A listed operator delegates when every operand casts to int64; an unlisted -# operator or an operand narrower than int32 falls back instead. +# operator or an operand narrower than int32 falls back instead. to_edge widens +# the operands of some operators, such as cat, before partitioning. int64_extension_partition_suite = { "add": (lambda x, y: x + y, (_WIDE, _WIDE_2), 1), "pad": (lambda x: torch.nn.functional.pad(x, (1, 1)), (_WIDE,), 1), "cumsum": (lambda x: torch.cumsum(x, 0), (_WIDE,), 0), - "cat_narrow": (lambda x, y: torch.cat([x, y]), (_WIDE, _NARROW), 1), + "cat_narrow_promoted": (lambda x, y: torch.cat([x, y]), (_WIDE, _NARROW), 1), "or_narrow": (torch.bitwise_or, (_WIDE, _NARROW), 1), - "cat_byte": (lambda x, y: torch.cat([x, y]), (_WIDE, _BYTES), 0), + "cat_byte_promoted": (lambda x, y: torch.cat([x, y]), (_WIDE, _BYTES), 1), + "or_byte": (torch.bitwise_or, (_WIDE, _BYTES), 0), "eq_bool": (lambda x, y: x == y, (_WIDE, _FLAGS), 0), - "cat_bool": (lambda x, y: torch.cat([x, y]), (_FLAGS, _WIDE), 1), + "cat_bool_promoted": (lambda x, y: torch.cat([x, y]), (_FLAGS, _WIDE), 1), + "or_bool": (torch.bitwise_or, (_FLAGS, _WIDE), 1), "where": (torch.where, (_COND, _WIDE, _WIDE_2), 0), "where_buffer": (Int64BufferWhere(), (_COND, _WIDE_2), 0), "strided_slice": (lambda x: x[::2], (torch.arange(5, dtype=torch.int64),), 0), diff --git a/backends/arm/test/misc/test_tosa_operator_support.py b/backends/arm/test/misc/test_tosa_operator_support.py index 25c40d63905..a65ec268355 100644 --- a/backends/arm/test/misc/test_tosa_operator_support.py +++ b/backends/arm/test/misc/test_tosa_operator_support.py @@ -334,3 +334,16 @@ def test_int64_operand_check_skips_symbolic_inputs() -> None: node = graph.call_function(exir_ops.edge.aten.view_copy.default, (x, [size, -1])) assert _int64_operands_agree(node) + + +@pytest.mark.parametrize( + ("dtype", "agree"), + ((torch.int64, True), (torch.int32, True), (torch.bool, True), (torch.int8, False)), +) +def test_int64_operand_check_widths(dtype, agree) -> None: + graph = torch.fx.Graph() + x = _placeholder(graph, "x", (4,), torch.int64) + y = _placeholder(graph, "y", (4,), dtype) + node = graph.call_function(exir_ops.edge.aten.bitwise_or.Tensor, (x, y)) + + assert _int64_operands_agree(node) == agree