Arm backend: Improve dtype validation (pytorch#15871)

oscarandersson8218 · web-flow · commit d536d1890d02 · 2025-12-04T15:06:31.000+01:00
Improve dtype validation in node-visitors.

Signed-off-by: Oscar Andersson &lt;oscar.andersson@arm.com&gt;
diff --git a/backends/arm/operator_support/ethos_u55_support.py b/backends/arm/operator_support/ethos_u55_support.py
@@ -78,10 +78,12 @@ def __init__(self, reporter: WhyNoPartitionReporter):
 
     targeted_ops_i8_i16_i32 = [
         exir_ops.edge.aten.cat.default,
+        exir_ops.edge.aten.expand_copy.default,
         exir_ops.edge.aten.repeat.default,
         exir_ops.edge.aten.constant_pad_nd.default,
         exir_ops.edge.aten.view.default,
         exir_ops.edge.aten.permute.default,
+        exir_ops.edge.aten.permute_copy.default,
     ]
 
     target_ops_i8 = tuple(TableOps.included_ops())
diff --git a/backends/arm/operators/op_avg_pool2d.py b/backends/arm/operators/op_avg_pool2d.py
@@ -115,10 +115,13 @@ def define_node(
     ) -> None:
         validate_num_inputs(self.target, inputs, [3, 4, 5, 6, 7])
         validate_same_dtype(self.target, [inputs[0], output], ts)
+        supported_dtypes = [ts.DType.INT8, ts.DType.FP32]
+        if self.tosa_spec.support_extension("int16"):
+            supported_dtypes.append(ts.DType.INT16)
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.FP32],
+            supported_dtypes,
             output.tosa_spec,
         )
 
diff --git a/backends/arm/operators/op_cat.py b/backends/arm/operators/op_cat.py
@@ -14,6 +14,8 @@
 )
 from executorch.backends.arm.operators.operator_validation_utils import (
     validate_num_inputs,
+    validate_same_dtype,
+    validate_valid_dtype,
 )
 from executorch.backends.arm.tosa.mapping import TosaArg
 from torch.fx import Node
@@ -35,9 +37,19 @@ def define_node(
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
+        supported_dtypes = [ts.DType.BOOL, ts.DType.INT8, ts.DType.INT32, ts.DType.FP32]
+        if self.tosa_spec.support_extension("int16"):
+            supported_dtypes.append(ts.DType.INT16)
         validate_num_inputs(self.target, inputs, [1, 2])
+        input_tosa_args = [TosaArg(arg, output.tosa_spec) for arg in inputs[0].special]
+        validate_same_dtype(self.target, [*input_tosa_args, output], ts)
+        validate_valid_dtype(
+            self.target,
+            [*input_tosa_args, output],
+            supported_dtypes,
+            output.tosa_spec,
+        )
 
-        tensors = inputs[0].special
         dim = 0 if len(inputs) < 2 else inputs[1].number
         rank = len(output.shape)
         dim = (dim + rank) % rank
@@ -50,7 +62,7 @@ def define_node(
             node,
             tosa_graph,
             ts.Op.CONCAT,
-            [tensor.name for tensor in tensors],
+            [tensor.name for tensor in input_tosa_args],
             [output.name],
             attr,
         )
diff --git a/backends/arm/operators/op_clamp.py b/backends/arm/operators/op_clamp.py
@@ -87,20 +87,18 @@ def define_node(
     ) -> None:
         validate_num_inputs(self.target, inputs, [2, 3])
         validate_same_dtype(self.target, [inputs[0], output], ts)
+        supported_dtypes = [ts.DType.INT8, ts.DType.FP16, ts.DType.FP32]
+        if self.tosa_spec.support_extension("int16"):
+            supported_dtypes.append(ts.DType.INT16)
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [
-                ts.DType.INT8,
-                ts.DType.INT16,
-                ts.DType.FP16,
-                ts.DType.FP32,
-            ],
+            supported_dtypes,
             output.tosa_spec,
         )
 
         node_input_dtype = node.meta["val"].dtype
-        # NOTE: Quantization of the min/max arguments is handled by QuantizeClampArgumentsPass
+        # NOTE: Quantization of the min/max arguments is handled by QuantizeOperatorArguments
         min_val, max_val = self._get_min_max_arguments(node, node_input_dtype)
 
         attr = ts.TosaSerializerAttribute()
diff --git a/backends/arm/operators/op_eq.py b/backends/arm/operators/op_eq.py
@@ -47,7 +47,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT32, ts.DType.FP32],
             output.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)
diff --git a/backends/arm/operators/op_ge.py b/backends/arm/operators/op_ge.py
@@ -47,7 +47,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT32, ts.DType.FP32],
             output.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)
diff --git a/backends/arm/operators/op_gt.py b/backends/arm/operators/op_gt.py
@@ -47,7 +47,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT32, ts.DType.FP32],
             output.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)
diff --git a/backends/arm/operators/op_index_select.py b/backends/arm/operators/op_index_select.py
@@ -12,6 +12,11 @@
     NodeVisitor,
     register_node_visitor,
 )
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+    validate_same_dtype,
+    validate_valid_dtype,
+)
 from executorch.backends.arm.tosa.mapping import TosaArg
 
 from executorch.backends.arm.tosa.utils import build_reshape_tosa_1_0
@@ -45,10 +50,16 @@ def define_node(
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
-        if len(inputs) != 3:
-            raise ValueError(f"Number of inputs are not 3: {len(inputs)}")
+        validate_num_inputs(self.target, inputs, 3)
+        validate_same_dtype(self.target, [inputs[0], output], ts)
+        validate_valid_dtype(
+            self.target,
+            [inputs[0], output],
+            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
+            output.tosa_spec,
+        )
 
-        weights, index, indices = inputs
+        weights, _, indices = inputs
 
         if len(weights.shape) == 2:
             weights_new_shape = [1, weights.shape[0], weights.shape[1]]
diff --git a/backends/arm/operators/op_le.py b/backends/arm/operators/op_le.py
@@ -47,7 +47,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT32, ts.DType.FP32],
             output.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)
diff --git a/backends/arm/operators/op_lt.py b/backends/arm/operators/op_lt.py
@@ -47,7 +47,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             inputs,
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT32, ts.DType.FP32],
             output.tosa_spec,
         )
         validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)
diff --git a/backends/arm/operators/op_max_pool2d.py b/backends/arm/operators/op_max_pool2d.py
@@ -44,10 +44,13 @@ def define_node(
     ) -> None:
         validate_num_inputs(self.target, inputs, [3, 4, 5, 6])
         validate_same_dtype(self.target, [inputs[0], output], ts)
+        supported_dtypes = [ts.DType.INT8, ts.DType.FP32]
+        if self.tosa_spec.support_extension("int16"):
+            supported_dtypes.append(ts.DType.INT16)
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.FP32],
+            supported_dtypes,
             output.tosa_spec,
         )
 
diff --git a/backends/arm/operators/op_mul.py b/backends/arm/operators/op_mul.py
@@ -44,7 +44,7 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [*inputs, output],
-            [ts.DType.INT32, ts.DType.FP32],
+            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
             output.tosa_spec,
         )
 
diff --git a/backends/arm/operators/op_permute.py b/backends/arm/operators/op_permute.py
@@ -116,7 +116,13 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
+            [
+                ts.DType.BOOL,
+                ts.DType.INT8,
+                ts.DType.INT16,
+                ts.DType.INT32,
+                ts.DType.FP32,
+            ],
             output.tosa_spec,
         )
 
diff --git a/backends/arm/operators/op_repeat.py b/backends/arm/operators/op_repeat.py
@@ -43,7 +43,13 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [ts.DType.INT8, ts.DType.INT32, ts.DType.INT16, ts.DType.FP32],
+            [
+                ts.DType.BOOL,
+                ts.DType.INT8,
+                ts.DType.INT16,
+                ts.DType.INT32,
+                ts.DType.FP32,
+            ],
             output.tosa_spec,
         )
 
diff --git a/backends/arm/operators/op_slice.py b/backends/arm/operators/op_slice.py
@@ -73,7 +73,13 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],
+            [
+                ts.DType.BOOL,
+                ts.DType.INT8,
+                ts.DType.INT16,
+                ts.DType.INT32,
+                ts.DType.FP32,
+            ],
             output.tosa_spec,
         )
 
diff --git a/backends/arm/operators/op_sum.py b/backends/arm/operators/op_sum.py
@@ -15,6 +15,7 @@
 from executorch.backends.arm.operators.operator_validation_utils import (
     validate_num_inputs,
     validate_same_dtype,
+    validate_valid_dtype,
 )
 from executorch.backends.arm.tosa import TosaSpecification
 from executorch.backends.arm.tosa.mapping import TosaArg
@@ -39,6 +40,12 @@ def define_node(
     ) -> None:
         validate_num_inputs(self.target, inputs, 3)
         validate_same_dtype(self.target, [inputs[0], output], ts)
+        validate_valid_dtype(
+            self.target,
+            [inputs[0], output],
+            [ts.DType.INT32, ts.DType.FP32],
+            output.tosa_spec,
+        )
 
         tensor = inputs[0]
         input_shape = list(tensor.shape)
diff --git a/backends/arm/operators/op_tosa_conv2d.py b/backends/arm/operators/op_tosa_conv2d.py
@@ -23,7 +23,7 @@
     validate_valid_dtype,
 )
 from executorch.backends.arm.tosa.mapping import TosaArg
-from executorch.backends.arm.tosa.specification import Tosa_1_00, TosaSpecification
+from executorch.backends.arm.tosa.specification import TosaSpecification
 
 
 @register_node_visitor
@@ -64,9 +64,7 @@ def define_node(
         if self.tosa_spec.support_integer():
             valid_input_dtypes.append(ts.DType.INT8)
 
-        if isinstance(self.tosa_spec, Tosa_1_00) and self.tosa_spec.support_extension(
-            "int16"
-        ):
+        if self.tosa_spec.support_extension("int16"):
             valid_input_dtypes.append(ts.DType.INT16)
             # Check constraints for int16 activations
             if inputs[0].dtype == ts.DType.INT16:
diff --git a/backends/arm/operators/op_tosa_matmul.py b/backends/arm/operators/op_tosa_matmul.py
@@ -51,16 +51,22 @@ def define_node(
         """Define the TOSA ``MATMUL`` operator."""
         validate_num_inputs(self.target, inputs, 2)
         validate_same_dtype(self.target, [*inputs], ts)
+        supported_input_dtypes = [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32]
+        if self.tosa_spec.support_extension("int16"):
+            supported_input_dtypes.append(ts.DType.INT16)
         validate_valid_dtype(
             self.target,
             [*inputs],
-            [ts.DType.INT8, ts.DType.INT16, ts.DType.FP32],
+            supported_input_dtypes,
             output.tosa_spec,
         )
+        supported_output_dtypes = [ts.DType.INT32, ts.DType.FP32]
+        if self.tosa_spec.support_extension("int16"):
+            supported_output_dtypes.append(ts.DType.INT48)
         validate_valid_dtype(
             self.target,
             [output],
-            [ts.DType.INT32, ts.DType.INT48, ts.DType.FP32],
+            supported_output_dtypes,
             output.tosa_spec,
         )
 
diff --git a/backends/arm/operators/op_tosa_resize.py b/backends/arm/operators/op_tosa_resize.py
@@ -39,33 +39,31 @@ def define_node(
         output: TosaArg,
     ) -> None:
         validate_num_inputs(self.target, inputs, [3, 4])
+        supported_input_dtypes = [ts.DType.INT8, ts.DType.FP32]
+        if self.tosa_spec.support_extension("int16"):
+            supported_input_dtypes.append(ts.DType.INT16)
+        validate_valid_dtype(
+            self.target,
+            [inputs[0]],
+            supported_input_dtypes,
+            output.tosa_spec,
+        )
+        supported_output_dtypes = [ts.DType.FP32]
         if node.kwargs.get("resize_mode") == "bilinear":
             resize_mode = ts.ResizeMode.BILINEAR
             align_corners = bool(node.args[2])
+            supported_output_dtypes.append(ts.DType.INT32)
+            if self.tosa_spec.support_extension("int16"):
+                supported_output_dtypes.append(ts.DType.INT48)
         else:
             resize_mode = ts.ResizeMode.NEAREST
             align_corners = False
             validate_same_dtype(self.target, [inputs[0], output], ts)
-
-        valid_dtypes = []
-        if self.tosa_spec.support_integer():
-            valid_dtypes.extend(
-                [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.INT48]
-            )
-
-        if self.tosa_spec.support_float():
-            valid_dtypes.extend(
-                [
-                    ts.DType.FP16,
-                    ts.DType.FP32,
-                ]
-            )
-
+            supported_output_dtypes.append(ts.DType.INT8)
+            if self.tosa_spec.support_extension("int16"):
+                supported_output_dtypes.append(ts.DType.INT16)
         validate_valid_dtype(
-            self.target,
-            [inputs[0], output],
-            valid_dtypes,
-            output.tosa_spec,
+            self.target, [output], supported_output_dtypes, output.tosa_spec
         )
         # tosa_shape output is NHWC, take HW
         input_size_yx = tuple([inputs[0].shape[dim] for dim in inputs[0].dim_order])[
diff --git a/backends/arm/operators/op_tosa_table.py b/backends/arm/operators/op_tosa_table.py
@@ -36,13 +36,18 @@ def define_node(
         output: TosaArg,
     ) -> None:
         validate_num_inputs(self.target, inputs, 2)
+        supported_input_dtypes = [ts.DType.INT8]
+        supported_output_dtypes = [ts.DType.INT8]
+        if self.tosa_spec.support_extension("int16"):
+            supported_input_dtypes.append(ts.DType.INT16)
+            supported_output_dtypes.append(ts.DType.INT32)
+
+        validate_valid_dtype(
+            self.target, inputs, supported_input_dtypes, output.tosa_spec
+        )
         validate_valid_dtype(
-            self.target, inputs, [ts.DType.INT8, ts.DType.INT16], output.tosa_spec
+            self.target, output, supported_output_dtypes, output.tosa_spec
         )
-        if inputs[0].dtype == ts.DType.INT8:
-            validate_valid_dtype(self.target, output, ts.DType.INT8, output.tosa_spec)
-        if inputs[0].dtype == ts.DType.INT16:
-            validate_valid_dtype(self.target, output, ts.DType.INT32, output.tosa_spec)
 
         # The name of the table constant is a bit complex.
         # The name of the pytorch buffer will be the target of last node argument.
diff --git a/backends/arm/operators/op_tosa_transpose.py b/backends/arm/operators/op_tosa_transpose.py
@@ -47,12 +47,12 @@ def define_node(
             self.target,
             [inputs[0], output],
             [
+                ts.DType.BOOL,
                 ts.DType.INT8,
                 ts.DType.INT16,
                 ts.DType.INT32,
-                ts.DType.FP32,
-                ts.DType.BOOL,
                 ts.DType.FP16,
+                ts.DType.FP32,
             ],
             output.tosa_spec,
         )
diff --git a/backends/arm/operators/op_where.py b/backends/arm/operators/op_where.py
diff --git a/backends/arm/operators/ops_identity.py b/backends/arm/operators/ops_identity.py
diff --git a/backends/arm/test/ops/test_expand.py b/backends/arm/test/ops/test_expand.py
diff --git a/backends/arm/test/ops/test_permute.py b/backends/arm/test/ops/test_permute.py
diff --git a/backends/arm/test/ops/test_repeat.py b/backends/arm/test/ops/test_repeat.py

Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ def define_node(`
`47`	`47`	`validate_valid_dtype(`
`48`	`48`	`self.target,`
`49`	`49`	`inputs,`
`50`		`- [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32],`
	`50`	`+ [ts.DType.INT32, ts.DType.FP32],`
`51`	`51`	`output.tosa_spec,`
`52`	`52`	`)`
`53`	`53`	`validate_valid_dtype(self.target, output, ts.DType.BOOL, output.tosa_spec)`