Arm backend: Create op utility function for num input verification (#10713)

Create function for validating that the number of inputs is correct in
each node visitor.

Signed-off-by: Sebastian Larsson <[email protected]>

Author: Sebastian-Larsson

backends/arm/operators/op_abs.py

@@ -13,6 +13,9 @@
     NodeVisitor,
     register_node_visitor,
 )
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+)
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from torch.fx import Node
@@ -39,6 +42,7 @@ def define_node(
 
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 1)
         # Specification (0.80) states that input and output types
         # should all be the same
         if not (inputs[0].dtype == output.dtype):
@@ -105,6 +109,7 @@ def define_node(
 
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 1)
         # Specification (0.80) states that input and output types
         # should all be the same
         if not (inputs[0].dtype == output.dtype):
@@ -157,6 +162,8 @@ def define_node(
 
         import serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 1)
+
         # Specification (1.0) states that input and output types
         # should all be the same
         if not (inputs[0].dtype == output.dtype):
@@ -224,6 +231,8 @@ def define_node(
 
         import serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 1)
+
         # Specification (1.0) states that input and output types
         # should all be the same
         if not (inputs[0].dtype == output.dtype):

backends/arm/operators/op_add.py

@@ -14,6 +14,9 @@
     NodeVisitor,
     register_node_visitor,
 )
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+)
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from torch.fx import Node
@@ -40,6 +43,7 @@ def define_node(
 
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 2)
         # Specification (0.80) states that input and output types
         # should all be the same
         if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
@@ -118,6 +122,7 @@ def define_node(
 
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 2)
         # Specification (0.80) states that input and output types
         # should all be the same
         if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
@@ -169,6 +174,8 @@ def define_node(
 
         import serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 2)
+
         # Specification (1.0) states that input and output types
         # should all be the same
         if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
@@ -237,6 +244,8 @@ def define_node(
 
         import serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 2)
+
         # Specification (1.0) states that input and output types
         # should all be the same
         if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:

backends/arm/operators/op_amax.py

@@ -9,6 +9,9 @@
     NodeVisitor,
     register_node_visitor,
 )
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+)
 from executorch.backends.arm.tosa_mapping import TosaArg
 from torch.fx import Node
 
@@ -31,6 +34,8 @@ def define_node(
     ) -> None:
         import tosa_tools.v0_80.serializer.tosa_serializer as ts
 
+        validate_num_inputs(self.target, inputs, 3)
+
         input = inputs[0]
         dim = inputs[1].number
 
@@ -71,6 +76,8 @@ def define_node(
     ) -> None:
         import serializer.tosa_serializer as ts
 
+        validate_num_inputs(self.target, inputs, 3)
+
         input = inputs[0]
         dim = inputs[1].number
 

backends/arm/operators/op_amin.py

@@ -9,6 +9,9 @@
     NodeVisitor,
     register_node_visitor,
 )
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+)
 from executorch.backends.arm.tosa_mapping import TosaArg
 from torch.fx import Node
 
@@ -31,6 +34,8 @@ def define_node(
     ) -> None:
         import tosa_tools.v0_80.serializer.tosa_serializer as ts
 
+        validate_num_inputs(self.target, inputs, 3)
+
         input = inputs[0]
         dim = inputs[1].number
 
@@ -71,6 +76,8 @@ def define_node(
     ) -> None:
         import serializer.tosa_serializer as ts
 
+        validate_num_inputs(self.target, inputs, 3)
+
         input = inputs[0]
         dim = inputs[1].number
 

backends/arm/operators/op_any.py

@@ -10,6 +10,9 @@
     NodeVisitor,
     register_node_visitor,
 )
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+)
 
 from executorch.backends.arm.tosa_mapping import TosaArg  # type: ignore
 from torch.fx import Node
@@ -30,6 +33,8 @@ def define_node(
     ) -> None:
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 3)
+
         if not (inputs[0].dtype == output.dtype):
             raise ValueError(
                 "All inputs and outputs need same dtype."
@@ -69,6 +74,8 @@ def define_node(
     ) -> None:
         import serializer.tosa_serializer as ts
 
+        validate_num_inputs(self.target, inputs, 3)
+
         if not (inputs[0].dtype == output.dtype):
             raise ValueError(
                 "All inputs and outputs need same dtype."

backends/arm/operators/op_avg_pool2d.py

@@ -16,6 +16,9 @@
     NodeVisitor,
     register_node_visitor,
 )
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+)
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_specification import TosaSpecification
 
@@ -85,6 +88,8 @@ def define_node(
     ) -> None:
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, [3, 4, 6])
+
         supported_dtypes = [ts.DType.INT8]
         if inputs[0].dtype not in supported_dtypes:
             raise TypeError(
@@ -122,6 +127,8 @@ def define_node(
     ) -> None:
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, [3, 4, 6])
+
         supported_dtypes = [ts.DType.INT8, ts.DType.FP32]
         if inputs[0].dtype not in supported_dtypes:
             raise TypeError(
@@ -212,6 +219,8 @@ def define_node(
     ) -> None:
         import serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, [3, 4, 6])
+
         supported_dtypes = [ts.DType.INT8]
         if inputs[0].dtype not in supported_dtypes:
             raise TypeError(
@@ -252,6 +261,8 @@ def define_node(
     ) -> None:
         import serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, [3, 4, 6])
+
         supported_dtypes = [ts.DType.INT8, ts.DType.FP32]
         if inputs[0].dtype not in supported_dtypes:
             raise TypeError(

backends/arm/operators/op_bmm.py

@@ -17,7 +17,9 @@
     NodeVisitor,
     register_node_visitor,
 )
-
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+)
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_quant_utils import build_rescale, build_rescale_v0_80
 from executorch.backends.arm.tosa_specification import TosaSpecification
@@ -46,6 +48,7 @@ def define_node(
 
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 2)
         if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
             raise TypeError(
                 f"All IO needs to have the same data type, got: "
@@ -128,6 +131,8 @@ def define_node(
 
         import serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 2)
+
         if inputs[0].dtype != inputs[1].dtype or inputs[0].dtype != output.dtype:
             raise TypeError(
                 f"All IO needs to have the same data type, got: "

backends/arm/operators/op_cat.py

@@ -11,6 +11,9 @@
     NodeVisitor,
     register_node_visitor,
 )
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+)
 from executorch.backends.arm.tosa_mapping import TosaArg
 from torch.fx import Node
 
@@ -33,6 +36,8 @@ def define_node(
     ) -> None:
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, [1, 2])
+
         tensors = inputs[0].special
         dim = 0 if len(inputs) < 2 else inputs[1].number
         rank = len(output.shape)
@@ -68,6 +73,8 @@ def define_node(
     ) -> None:
         import serializer.tosa_serializer as ts
 
+        validate_num_inputs(self.target, inputs, [1, 2])
+
         tensors = inputs[0].special
         dim = 0 if len(inputs) < 2 else inputs[1].number
         rank = len(output.shape)

backends/arm/operators/op_clamp.py

@@ -15,6 +15,9 @@
     NodeVisitor,
     register_node_visitor,
 )
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+)
 
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_specification import TosaSpecification
@@ -65,9 +68,6 @@ def cast_type(value: Any) -> int | float:
                 # Attempt to cast to float
                 return float(value)
 
-        if len(node.args) != 2 and len(node.args) != 3:
-            raise ValueError(f"Expected len(node.args) to be 2 or 3, got {node.args}")
-
         min_arg = dtype_min
         max_arg = dtype_max
 
@@ -87,10 +87,7 @@ def define_node(
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
-        if len(node.all_input_nodes) != 1:
-            raise ValueError(
-                f"Expected 1 input for {self.target}, got {len(node.all_input_nodes)}"
-            )
+        validate_num_inputs(self.target, inputs, [2, 3])
 
         min_int8, max_int8 = self._get_min_max_arguments(
             node,
@@ -130,10 +127,7 @@ def define_node(
     ) -> None:
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
-        if len(node.all_input_nodes) != 1:
-            raise ValueError(
-                f"Expected 1 input for {self.target}, got {len(node.all_input_nodes)}"
-            )
+        validate_num_inputs(self.target, inputs, [2, 3])
 
         if inputs[0].dtype == ts.DType.INT8:
             # Call the inherited define_node for handling integers
@@ -178,9 +172,6 @@ def cast_type(value: Any) -> int | float:
                 # Attempt to cast to float
                 return float(value)
 
-        if len(node.args) != 2 and len(node.args) != 3:
-            raise ValueError(f"Expected len(node.args) to be 2 or 3, got {node.args}")
-
         min_arg = dtype_min
         max_arg = dtype_max
 
@@ -202,10 +193,7 @@ def define_node(
     ) -> None:
         import serializer.tosa_serializer as ts  # type: ignore
 
-        if len(node.all_input_nodes) != 1:
-            raise ValueError(
-                f"Expected 1 input for {self.target}, got {len(node.all_input_nodes)}"
-            )
+        validate_num_inputs(self.target, inputs, [2, 3])
 
         # NOTE: Quantization of the min/max arguments is handled by QuantizeOperatorArguments
         min_int8, max_int8 = self._get_min_max_arguments(
@@ -247,10 +235,7 @@ def define_node(
     ) -> None:
         import serializer.tosa_serializer as ts  # type: ignore
 
-        if len(node.all_input_nodes) != 1:
-            raise ValueError(
-                f"Expected 1 input for {self.target}, got {len(node.all_input_nodes)}"
-            )
+        validate_num_inputs(self.target, inputs, [2, 3])
 
         min_fp32, max_fp32 = self._get_min_max_arguments(
             node,

backends/arm/operators/op_constant_pad_nd.py

@@ -16,6 +16,9 @@
     NodeVisitor,
     register_node_visitor,
 )
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+)
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_specification import TosaSpecification
 
@@ -39,6 +42,8 @@ def define_node(
     ) -> None:
         import tosa_tools.v0_80.serializer.tosa_serializer as ts
 
+        validate_num_inputs(self.target, inputs, 3)
+
         if inputs[0].dtype == ts.DType.INT8:
             input_qparams = get_input_qparams(node)
             qargs = input_qparams[0]
@@ -98,9 +103,10 @@ def define_node(
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
-
         import serializer.tosa_serializer as ts  # type: ignore
 
+        validate_num_inputs(self.target, inputs, 3)
+
         if inputs[0].dtype == ts.DType.INT8:
             input_qparams = get_input_qparams(node)
             qargs = input_qparams[0]