Arm backend: Make int64 representable under EXT-INT64 (1/3)

backends/arm/operators/op_flip.py

@@ -50,6 +50,9 @@ def define_node(
 
         validate_num_inputs(self.target, inputs, 2)
         validate_same_dtype(self.target, [inputs[0], output])
+        if self.tosa_spec.support_extension("int64"):
+            supported_dtypes.append(ts.DType.INT64)
+
         validate_valid_dtype(
             self.target,
             [inputs[0], output],

backends/arm/operators/op_tosa_abs.py

@@ -23,4 +23,5 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
             attr_method="AbsAttribute",
             num_inputs=1,
             input_dtypes=[ts.DType.INT32, ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
+            extension_dtypes={"int64": [ts.DType.INT64]},
         )

backends/arm/operators/op_tosa_add.py

@@ -23,4 +23,5 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
             attr_method="AddAttribute",
             num_inputs=2,
             input_dtypes=[ts.DType.INT32, ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
+            extension_dtypes={"int64": [ts.DType.INT64]},
         )

backends/arm/operators/op_tosa_binary_ops.py

@@ -19,6 +19,7 @@ def binary_operator_factory(
     tosa_op,
     attr_method: str,
     valid_dtypes: List[Any],
+    extension_dtypes: dict[str, List[Any]] | None = None,
 ):
     operator_target = target
 
@@ -32,6 +33,7 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
                 attr_method=attr_method,
                 num_inputs=2,
                 input_dtypes=valid_dtypes,
+                extension_dtypes=extension_dtypes,
             )
 
     register_node_visitor(BinaryOperator)
@@ -42,18 +44,21 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
     ts.Op.BITWISE_AND,
     "BitwiseAndAttribute",
     [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32],
+    {"int64": [ts.DType.INT64]},
 )
 binary_operator_factory(
     "tosa.BITWISE_OR.default",
     ts.Op.BITWISE_OR,
     "BitwiseOrAttribute",
     [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32],
+    {"int64": [ts.DType.INT64]},
 )
 binary_operator_factory(
     "tosa.BITWISE_XOR.default",
     ts.Op.BITWISE_XOR,
     "BitwiseXorAttribute",
     [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32],
+    {"int64": [ts.DType.INT64]},
 )
 binary_operator_factory(
     "tosa.INTDIV.default",

backends/arm/operators/op_tosa_bitwise_not.py

@@ -27,4 +27,5 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
             attr_method="BitwiseNotAttribute",
             num_inputs=1,
             input_dtypes=[ts.DType.INT8, ts.DType.INT16, ts.DType.INT32],
+            extension_dtypes={"int64": [ts.DType.INT64]},
         )

backends/arm/operators/op_tosa_eq.py

@@ -23,6 +23,7 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
             attr_method="EqualAttribute",
             num_inputs=2,
             input_dtypes=[ts.DType.INT32, ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
+            extension_dtypes={"int64": [ts.DType.INT64]},
             output_dtypes=[ts.DType.BOOL],
             same_dtype_with_output=False,
             dtype_check_inputs_only=True,

backends/arm/operators/op_tosa_ge.py

@@ -25,6 +25,7 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
             attr_method="GreaterEqualAttribute",
             num_inputs=2,
             input_dtypes=COMPARE_INPUT_DTYPES,
+            extension_dtypes={"int64": [ts.DType.INT64]},
             output_dtypes=[ts.DType.BOOL],
             same_dtype_with_output=False,
             dtype_check_inputs_only=True,

backends/arm/operators/op_tosa_gt.py

@@ -25,6 +25,7 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
             attr_method="GreaterAttribute",
             num_inputs=2,
             input_dtypes=COMPARE_INPUT_DTYPES,
+            extension_dtypes={"int64": [ts.DType.INT64]},
             output_dtypes=[ts.DType.BOOL],
             same_dtype_with_output=False,
             dtype_check_inputs_only=True,

backends/arm/operators/op_tosa_identity.py

@@ -35,4 +35,5 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
                 ts.DType.FP8E4M3,
                 ts.DType.FP8E5M2,
             ],
+            extension_dtypes={"int64": [ts.DType.INT64]},
         )

backends/arm/operators/op_tosa_maximum.py

@@ -24,4 +24,5 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
             attr_kwargs={"nan_mode": ts.NanPropagationMode.PROPAGATE},
             num_inputs=2,
             input_dtypes=[ts.DType.INT32, ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
+            extension_dtypes={"int64": [ts.DType.INT64]},
         )

backends/arm/operators/op_tosa_minimum.py

@@ -24,4 +24,5 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
             attr_kwargs={"nan_mode": ts.NanPropagationMode.PROPAGATE},
             num_inputs=2,
             input_dtypes=[ts.DType.INT32, ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
+            extension_dtypes={"int64": [ts.DType.INT64]},
         )

backends/arm/operators/op_tosa_mul.py

@@ -33,18 +33,18 @@ def define_node(
     ) -> None:
         validate_num_inputs(self.target, inputs, 2)
         validate_same_dtype(self.target, [*inputs, output])
-        validate_valid_dtype(
-            self.target,
-            [*inputs, output],
-            [
-                ts.DType.INT8,
-                ts.DType.INT16,
-                ts.DType.INT32,
-                ts.DType.FP32,
-                ts.DType.BF16,
-                ts.DType.FP16,
-            ],
-        )
+        supported_dtypes = [
+            ts.DType.INT8,
+            ts.DType.INT16,
+            ts.DType.INT32,
+            ts.DType.FP32,
+            ts.DType.BF16,
+            ts.DType.FP16,
+        ]
+        if self.tosa_spec.support_extension("int64"):
+            supported_dtypes.append(ts.DType.INT64)
+
+        validate_valid_dtype(self.target, [*inputs, output], supported_dtypes)
 
         shift = tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{output.name}_shift")
         attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_tosa_neg.py

@@ -55,6 +55,9 @@ def define_node(
 
         validate_num_inputs(self.target, inputs, 1)
         validate_same_dtype(self.target, [*inputs, output])
+        if self.tosa_spec.support_extension("int64"):
+            supported_dtypes.append(ts.DType.INT64)
+
         validate_valid_dtype(self.target, [*inputs, output], supported_dtypes)
 
         input_zp, output_zp = get_negate_zero_points(

backends/arm/operators/op_tosa_sub.py

@@ -23,4 +23,5 @@ def get_config(cls) -> SimpleNodeVisitorConfig:
             attr_method="SubAttribute",
             num_inputs=2,
             input_dtypes=[ts.DType.INT32, ts.DType.FP16, ts.DType.FP32, ts.DType.BF16],
+            extension_dtypes={"int64": [ts.DType.INT64]},
         )

backends/arm/operators/op_where.py

@@ -50,6 +50,9 @@ def define_node(
         # Not first input, which is condition tensor.
         validate_same_dtype(self.target, inputs[1:])
         validate_valid_dtype(self.target, inputs[0], ts.DType.BOOL)
+        if self.tosa_spec.support_extension("int64"):
+            supported_dtypes.append(ts.DType.INT64)
+
         validate_valid_dtype(
             self.target,
             [*inputs[1:], output],

backends/arm/operators/operator_validation_utils.py

@@ -37,12 +37,14 @@ def supported_data_layout_dtypes(
                 supported_dtypes.append(ts.DType.INT16)
     if tosa_spec.support_float():
         supported_dtypes.extend([ts.DType.FP16, ts.DType.FP32])
-    if tosa_spec.support_extension("bf16"):
-        supported_dtypes.append(ts.DType.BF16)
-    if tosa_spec.support_extension("fp8e4m3"):
-        supported_dtypes.append(ts.DType.FP8E4M3)
-    if tosa_spec.support_extension("fp8e5m2"):
-        supported_dtypes.append(ts.DType.FP8E5M2)
+    for extension, dtype in (
+        ("int64", ts.DType.INT64),
+        ("bf16", ts.DType.BF16),
+        ("fp8e4m3", ts.DType.FP8E4M3),
+        ("fp8e5m2", ts.DType.FP8E5M2),
+    ):
+        if tosa_spec.support_extension(extension):
+            supported_dtypes.append(dtype)
     if include_mxfp and tosa_spec.support_extension("mxfp"):
         for dtype in (ts.DType.FP8E4M3, ts.DType.FP8E5M2):
             if dtype not in supported_dtypes:

backends/arm/operators/simple_node_visitor.py

@@ -26,6 +26,7 @@ class SimpleNodeVisitorConfig:
     output_dtypes: List[Any] | None = None
     same_dtype_with_output: bool = True
     dtype_check_inputs_only: bool = False
+    extension_dtypes: dict[str, List[Any]] | None = None
 
 
 class SimpleNodeVisitor(NodeVisitor):
@@ -48,12 +49,17 @@ def define_node(
     ) -> None:
         cfg = self.get_config()
 
+        input_dtypes = list(cfg.input_dtypes)
+        for extension, dtypes in (cfg.extension_dtypes or {}).items():
+            if self.tosa_spec.support_extension(extension):
+                input_dtypes.extend(dtypes)
+
         self.validate(
             target=self.target,
             inputs=inputs,
             output=output,
             num_inputs=cfg.num_inputs,
-            input_dtypes=cfg.input_dtypes,
+            input_dtypes=input_dtypes,
             output_dtypes=cfg.output_dtypes,
             same_dtype_with_output=cfg.same_dtype_with_output,
             dtype_check_inputs_only=cfg.dtype_check_inputs_only,

backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py

@@ -385,3 +385,41 @@ def test_extrema_ops_reject_int32_on_fp_profile(op_name: str) -> None:
             getattr(exir_ops.backend.tosa, op_name).default(
                 *_to_fake(mode, input1, input2)
             )
+
+
+_INT64_BINARY_OPS = [
+    "ADD",
+    "SUB",
+    "MUL",
+    "MAXIMUM",
+    "MINIMUM",
+    "BITWISE_OR",
+    "BITWISE_XOR",
+]
+
+
+@pytest.mark.parametrize("op_name", _INT64_BINARY_OPS)
+@pytest.mark.parametrize(
+    ("spec", "accepted"), [("TOSA-1.1+INT+int64", True), ("TOSA-1.1+INT", False)]
+)
+def test_binary_ops_take_int64_only_with_the_extension(
+    op_name: str, spec: str, accepted: bool
+) -> None:
+    input1 = torch.randint(1, 16, (2, 1, 3), dtype=torch.int64)
+    input2 = torch.randint(1, 8, (1, 4, 3), dtype=torch.int64)
+
+    with TosaLoweringContext(
+        TosaSpecification.create_from_string(spec)
+    ), FakeTensorMode() as mode:
+        if not accepted:
+            with pytest.raises(TosaValueError):
+                getattr(exir_ops.backend.tosa, op_name).default(
+                    *_to_fake(mode, input1, input2)
+                )
+            return
+        output = getattr(exir_ops.backend.tosa, op_name).default(
+            *_to_fake(mode, input1, input2)
+        )
+
+    assert output.dtype == torch.int64
+    assert tuple(output.shape) == (2, 4, 3)

backends/arm/test/misc/tosa_dialect/test_tosa_dialect_comparison.py

@@ -95,3 +95,30 @@ def test_compare_ops_reject_int32_on_fp_profile(op_name: str) -> None:
             getattr(exir_ops.backend.tosa, op_name).default(
                 *_to_fake(mode, input1, input2)
             )
+
+
+@pytest.mark.parametrize("op_name", ["EQUAL", "GREATER", "GREATER_EQUAL"])
+@pytest.mark.parametrize(
+    ("spec", "accepted"), [("TOSA-1.1+INT+int64", True), ("TOSA-1.1+INT", False)]
+)
+def test_compare_ops_take_int64_only_with_the_extension(
+    op_name: str, spec: str, accepted: bool
+) -> None:
+    input1 = torch.randint(1, 16, (2, 1, 3), dtype=torch.int64)
+    input2 = torch.randint(1, 8, (1, 4, 3), dtype=torch.int64)
+
+    with TosaLoweringContext(
+        TosaSpecification.create_from_string(spec)
+    ), FakeTensorMode() as mode:
+        if not accepted:
+            with pytest.raises(TosaValueError):
+                getattr(exir_ops.backend.tosa, op_name).default(
+                    *_to_fake(mode, input1, input2)
+                )
+            return
+        output = getattr(exir_ops.backend.tosa, op_name).default(
+            *_to_fake(mode, input1, input2)
+        )
+
+    assert output.dtype == torch.bool
+    assert tuple(output.shape) == (2, 4, 3)

backends/arm/test/misc/tosa_dialect/test_tosa_dialect_ternary_ops.py

@@ -147,3 +147,28 @@ def test_tosa_select_rejects_bfloat16_without_bf16_extension() -> None:
             exir_ops.backend.tosa.SELECT.default(
                 *_to_fake(mode, condition, input1, input2)
             )
+
+
+@pytest.mark.parametrize(
+    ("spec", "accepted"), [("TOSA-1.1+INT+int64", True), ("TOSA-1.1+INT", False)]
+)
+def test_select_takes_int64_only_with_the_extension(spec: str, accepted: bool) -> None:
+    condition = torch.randint(0, 2, (2, 1, 3), dtype=torch.bool)
+    input1 = torch.randint(-8, 8, (1, 4, 3), dtype=torch.int64)
+    input2 = torch.randint(-8, 8, (1, 4, 3), dtype=torch.int64)
+
+    with TosaLoweringContext(
+        TosaSpecification.create_from_string(spec)
+    ), FakeTensorMode() as mode:
+        if not accepted:
+            with pytest.raises(TosaValueError):
+                exir_ops.backend.tosa.SELECT.default(
+                    *_to_fake(mode, condition, input1, input2)
+                )
+            return
+        output = exir_ops.backend.tosa.SELECT.default(
+            *_to_fake(mode, condition, input1, input2)
+        )
+
+    assert output.dtype == torch.int64
+    assert tuple(output.shape) == (2, 4, 3)

backends/arm/test/misc/tosa_dialect/test_tosa_dialect_unary_ops.py

@@ -392,3 +392,29 @@ def test_logical_not_rejects_non_bool() -> None:
     ), FakeTensorMode() as mode:
         with pytest.raises(TosaValueError, match="requires bool inputs"):
             exir_ops.backend.tosa.LOGICAL_NOT.default(mode.from_tensor(sample_input))
+
+
+@pytest.mark.parametrize("op_name", ["ABS", "NEGATE", "BITWISE_NOT"])
+@pytest.mark.parametrize(
+    ("spec", "accepted"), [("TOSA-1.1+INT+int64", True), ("TOSA-1.1+INT", False)]
+)
+def test_unary_ops_take_int64_only_with_the_extension(
+    op_name: str, spec: str, accepted: bool
+) -> None:
+    input_tensor = torch.randint(1, 16, (2, 3), dtype=torch.int64)
+
+    with TosaLoweringContext(
+        TosaSpecification.create_from_string(spec)
+    ), FakeTensorMode() as mode:
+        if not accepted:
+            with pytest.raises(TosaValueError):
+                getattr(exir_ops.backend.tosa, op_name).default(
+                    mode.from_tensor(input_tensor)
+                )
+            return
+        output = getattr(exir_ops.backend.tosa, op_name).default(
+            mode.from_tensor(input_tensor)
+        )
+
+    assert output.dtype == torch.int64
+    assert tuple(output.shape) == (2, 3)

backends/arm/tosa/dialect/ops/_common.py

@@ -16,6 +16,19 @@
 _VALID_NAN_MODES = {"PROPAGATE", "IGNORE"}
 
 
+def accepts_int64(dtype: torch.dtype, op: str) -> bool:
+    """Whether `dtype` is an int64 the active specification carries."""
+    if dtype != torch.int64:
+        return False
+    if not get_context_spec().support_extension("int64"):
+        raise TosaValueError(
+            f"TOSA spec {get_context_spec()} doesn't support "
+            f"{str(dtype).removeprefix('torch.')} for {op}",
+            op=op,
+        )
+    return True
+
+
 def broadcast_shape(
     input1: torch.Tensor, input2: torch.Tensor, op: str
 ) -> tuple[int | torch.SymInt, ...]: