Fix type-checking issues in Arm backend

Signed-off-by: Erik Lundell <[email protected]>
Change-Id: I4feb4b5c6d269d7c0ff4312c17fec52da413fa5a

Author: Erik-Lundell

backends/arm/operators/op_addmm.py

@@ -16,8 +16,8 @@
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_quant_utils import (
     build_rescale,
-    search_quant_arg_downstream,
-    search_quant_arg_upstream,
+    get_quant_arg_downstream,
+    get_quant_arg_upstream,
 )
 
 from executorch.backends.arm.tosa_utils import build_reshape
@@ -70,7 +70,7 @@ def define_node(
         input_zp = 0
         if is_quant_node:
             input_node = node.all_input_nodes[1]
-            input_zp = search_quant_arg_upstream(input_node).zp
+            input_zp = get_quant_arg_upstream(input_node).zp
         attr.ConvAttribute(
             pad=pad_attr,
             stride=stride_attr,
@@ -105,16 +105,16 @@ def define_node(
             # Read inputs' parent nodes
             _, input_node, weight_node = node.all_input_nodes
 
-            qargs = search_quant_arg_upstream(input_node)
+            qargs = get_quant_arg_upstream(input_node)
             input_scale = qargs.scale
             consumer_node = list(node.users)[0]
-            quant_args = search_quant_arg_downstream(consumer_node)
+            quant_args = get_quant_arg_downstream(consumer_node)
 
             consumer_node_scale = quant_args.scale
             consumer_node_node_zp = quant_args.zp
 
             weight_node_q_node = weight_node.all_input_nodes[0]
-            weight_scale = search_quant_arg_upstream(weight_node_q_node).scale
+            weight_scale = get_quant_arg_upstream(weight_node_q_node).scale
 
             output_rescale_scale = (input_scale * weight_scale) / consumer_node_scale
 

backends/arm/operators/op_bmm.py

@@ -16,8 +16,8 @@
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_quant_utils import (
     build_rescale,
-    search_quant_arg_downstream,
-    search_quant_arg_upstream,
+    get_quant_arg_downstream,
+    get_quant_arg_upstream,
 )
 from executorch.backends.arm.tosa_utils import get_two_inputs
 from serializer.tosa_serializer import TosaOp
@@ -46,8 +46,8 @@ def define_node(
         # For INT8, we need to get the zero points and add an intermediate tensor
         # for a later rescale.
         if is_quant_node:
-            input0_q_params = search_quant_arg_upstream(input0)
-            input1_q_params = search_quant_arg_upstream(input1)
+            input0_q_params = get_quant_arg_upstream(input0)
+            input1_q_params = get_quant_arg_upstream(input1)
             input0_zp = input0_q_params.zp
             input1_zp = input1_q_params.zp
             bmm_result = tosa_graph.addIntermediate(output.shape, ts.DType.INT32)
@@ -69,7 +69,7 @@ def define_node(
 
         # As INT8 accumulates into INT32, we need to rescale it back to INT8
         if is_quant_node:
-            output_q_params = search_quant_arg_downstream(list(node.users)[0])
+            output_q_params = get_quant_arg_downstream(list(node.users)[0])
 
             final_output_scale = (
                 input0_q_params.scale * input1_q_params.scale

backends/arm/operators/op_conv2d.py

@@ -15,8 +15,8 @@
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_quant_utils import (
     build_rescale_conv_output,
-    search_quant_arg_downstream,
-    search_quant_arg_upstream,
+    get_quant_arg_downstream,
+    get_quant_arg_upstream,
 )
 from executorch.backends.arm.tosa_utils import build_reshape, tosa_shape
 
@@ -83,9 +83,7 @@ def define_node(
         )
 
         input_zp = (
-            search_quant_arg_upstream(node.all_input_nodes[0]).zp
-            if is_quant_node
-            else 0
+            get_quant_arg_upstream(node.all_input_nodes[0]).zp if is_quant_node else 0
         )
 
         attr.ConvAttribute(
@@ -161,9 +159,9 @@ def define_node(
         # integer value domain of the next op. Otherwise return float32 output.
         if is_quant_node:
             # Get scale_factor from input, weight, and output.
-            input_scale = search_quant_arg_upstream(node.all_input_nodes[0]).scale
-            weight_scale = search_quant_arg_upstream(node.all_input_nodes[1]).scale
-            output_qargs = search_quant_arg_downstream(list(node.users)[0])
+            input_scale = get_quant_arg_upstream(node.all_input_nodes[0]).scale
+            weight_scale = get_quant_arg_upstream(node.all_input_nodes[1]).scale
+            output_qargs = get_quant_arg_downstream(list(node.users)[0])
 
             build_rescale_conv_output(
                 tosa_graph,

backends/arm/operators/op_exp.py

@@ -17,10 +17,10 @@
 
 from executorch.backends.arm.tosa_quant_utils import (
     dequantize_value,
+    get_quant_arg_downstream,
+    get_quant_arg_upstream,
     QuantArgs,
     quantize_value,
-    search_quant_arg_downstream,
-    search_quant_arg_upstream,
 )
 from serializer.tosa_serializer import TosaOp
 from torch.fx import Node
@@ -49,9 +49,9 @@ def define_node(
 
             # Create attribute for 8 bit table lookup.
             input_node = node.all_input_nodes[0]
-            in_quantargs = search_quant_arg_upstream(input_node)
+            in_quantargs = get_quant_arg_upstream(input_node)
             output_node = list(node.users)[0]
-            out_quantargs = search_quant_arg_downstream(output_node)
+            out_quantargs = get_quant_arg_downstream(output_node)
 
             table = exp_table_8bit(in_quantargs, out_quantargs)
             table_attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_full.py

@@ -15,8 +15,8 @@
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_quant_utils import (
+    get_quant_arg_downstream,
     quantize_value,
-    search_quant_arg_downstream,
 )
 from executorch.backends.arm.tosa_utils import tosa_shape
 from torch.fx import Node
@@ -42,7 +42,7 @@ def define_node(
 
         value = inputs[1].number
         if is_quant_node:
-            qargs = search_quant_arg_downstream(list(node.users)[0])
+            qargs = get_quant_arg_downstream(list(node.users)[0])
             qvalue = quantize_value(value, qargs)
             dtype = ts.DType.INT8
             data = np.full(shape, qvalue, dtype=np.int8)

backends/arm/operators/op_hardtanh.py

@@ -15,8 +15,8 @@
 from executorch.backends.arm.tosa_mapping import TosaArg
 
 from executorch.backends.arm.tosa_quant_utils import (
+    get_quant_arg_upstream,
     quantize_value,
-    search_quant_arg_upstream,
 )
 from serializer.tosa_serializer import TosaOp
 
@@ -40,7 +40,7 @@ def define_node(
 
         if is_quant_node:
             # Get quant parameters
-            qargs = search_quant_arg_upstream(node.all_input_nodes[0])
+            qargs = get_quant_arg_upstream(node.all_input_nodes[0])
             # Convert to quantized representation
             clamp_min_qs = quantize_value(inputs[1].number, qargs)
             clamp_max_qs = quantize_value(inputs[2].number, qargs)

backends/arm/operators/op_log.py

@@ -17,10 +17,10 @@
 
 from executorch.backends.arm.tosa_quant_utils import (
     dequantize_value,
+    get_quant_arg_downstream,
+    get_quant_arg_upstream,
     QuantArgs,
     quantize_value,
-    search_quant_arg_downstream,
-    search_quant_arg_upstream,
 )
 from serializer.tosa_serializer import TosaOp
 from torch.fx import Node
@@ -50,9 +50,9 @@ def define_node(
 
             # Create attribute for 8 bit table lookup.
             input_node = node.all_input_nodes[0]
-            in_quantargs = search_quant_arg_upstream(input_node)
+            in_quantargs = get_quant_arg_upstream(input_node)
             output_node = list(node.users)[0]
-            out_quantargs = search_quant_arg_downstream(output_node)
+            out_quantargs = get_quant_arg_downstream(output_node)
 
             table = log_table_8bit(in_quantargs, out_quantargs)
             table_attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_max_pool2d.py

@@ -14,8 +14,8 @@
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_utils import (
-    search_quant_arg_downstream,
-    search_quant_arg_upstream,
+    get_quant_arg_downstream,
+    get_quant_arg_upstream,
 )
 
 from serializer.tosa_serializer import TosaOp
@@ -57,10 +57,8 @@ def define_node(
         output_zp = 0
 
         if is_quant_node:
-            input_zp = search_quant_arg_upstream(
-                torch.fx.Node, node.all_input_nodes[0]
-            ).zp
-            output_zp = search_quant_arg_downstream(list(node.users)[0]).zp
+            input_zp = get_quant_arg_upstream(torch.fx.Node, node.all_input_nodes[0]).zp
+            output_zp = get_quant_arg_downstream(list(node.users)[0]).zp
 
         attr = ts.TosaSerializerAttribute()
         attr.PoolAttribute(

backends/arm/operators/op_mm.py

@@ -16,8 +16,8 @@
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_quant_utils import (
     build_rescale,
-    search_quant_arg_downstream,
-    search_quant_arg_upstream,
+    get_quant_arg_downstream,
+    get_quant_arg_upstream,
 )
 from executorch.backends.arm.tosa_utils import (
     build_reshape,
@@ -58,8 +58,8 @@ def define_node(
         # For INT8, we need to get the zero point, otherwise it is 0
         input0_zp, input1_zp = 0, 0
         if is_quant_node:
-            input0_zp = search_quant_arg_upstream(input0).zp
-            input1_zp = search_quant_arg_upstream(input1).zp
+            input0_zp = get_quant_arg_upstream(input0).zp
+            input1_zp = get_quant_arg_upstream(input1).zp
 
         mat_mul_result = tosa_graph.addIntermediate(
             output_new_shape, ts.DType.INT32 if is_quant_node else output.dtype
@@ -90,9 +90,9 @@ def define_node(
 
         # As INT8 accumulates into INT32, we need to rescale it back to INT8
         if is_quant_node:
-            input0_q_params = search_quant_arg_upstream(input0)
-            input1_q_params = search_quant_arg_upstream(input1)
-            output_q_params = search_quant_arg_downstream(list(node.users)[0])
+            input0_q_params = get_quant_arg_upstream(input0)
+            input1_q_params = get_quant_arg_upstream(input1)
+            output_q_params = get_quant_arg_downstream(list(node.users)[0])
 
             final_output_scale = (
                 input0_q_params.scale * input1_q_params.scale

backends/arm/operators/op_mul.py

@@ -37,10 +37,10 @@ def define_node(
         if is_quant_node:
             input_A = inputs[0]
             input_B = inputs[1]
-            input_A_qargs = tqutils.search_quant_arg_upstream(
+            input_A_qargs = tqutils.get_quant_arg_upstream(
                 cast(torch.fx.Node, node.args[0])
             )
-            input_B_qargs = tqutils.search_quant_arg_upstream(
+            input_B_qargs = tqutils.get_quant_arg_upstream(
                 cast(torch.fx.Node, node.args[1])
             )
 

backends/arm/operators/op_placeholder.py

@@ -10,9 +10,9 @@
 import torch.fx
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_quant_utils import (
+    get_quant_arg_upstream,
     get_quantized_node_output_dtype,
     is_node_quantized,
-    search_quant_arg_upstream,
 )
 from executorch.backends.arm.tosa_specification import TosaSpecification
 from executorch.backends.arm.tosa_utils import (
@@ -82,8 +82,8 @@ def process_quantized_bias(
             _,
         ) = consumer_node.all_input_nodes
 
-    input_node_scale = search_quant_arg_upstream(input_node).scale
-    weight_node_scale = search_quant_arg_upstream(weight_node).scale
+    input_node_scale = get_quant_arg_upstream(input_node).scale
+    weight_node_scale = get_quant_arg_upstream(weight_node).scale
     bias_values_quantized = (
         (parameter_values / (input_node_scale * weight_node_scale))
         .round()

backends/arm/operators/op_reciprocal.py

@@ -15,10 +15,10 @@
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_quant_utils import (
     dequantize_value,
+    get_quant_arg_downstream,
+    get_quant_arg_upstream,
     QuantArgs,
     quantize_value,
-    search_quant_arg_downstream,
-    search_quant_arg_upstream,
 )
 from serializer.tosa_serializer import TosaOp
 
@@ -42,8 +42,8 @@ def define_node(
 
         if is_quant_node:
             input = inputs[0]
-            input_qargs = search_quant_arg_upstream(node.all_input_nodes[0])
-            output_qargs = search_quant_arg_downstream(list(node.users)[0])
+            input_qargs = get_quant_arg_upstream(node.all_input_nodes[0])
+            output_qargs = get_quant_arg_downstream(list(node.users)[0])
 
             div_table = div_table_8bit(input_qargs, output_qargs)
 

backends/arm/operators/op_relu.py

@@ -38,7 +38,7 @@ def define_node(
         clamp_min_qs = 0
         clamp_max_qs = 0
         if is_quant_node:
-            out_qargs = tqutils.search_quant_arg_downstream(list(node.users)[0])
+            out_qargs = tqutils.get_quant_arg_downstream(list(node.users)[0])
             clamp_min_qs = tqutils.quantize_value(0, out_qargs)
             clamp_max_qs = tqutils.quantize_value(float("inf"), out_qargs)
 

backends/arm/operators/op_rsqrt.py

@@ -16,10 +16,10 @@
 from executorch.backends.arm.tosa_mapping import TosaArg
 from executorch.backends.arm.tosa_quant_utils import (
     dequantize_value,
+    get_quant_arg_downstream,
+    get_quant_arg_upstream,
     QuantArgs,
     quantize_value,
-    search_quant_arg_downstream,
-    search_quant_arg_upstream,
 )
 from serializer.tosa_serializer import TosaOp
 
@@ -40,9 +40,9 @@ def define_node(
             # Assume quantized input is 8 bit.
             # Create attribute for 8 bit table lookup.
             input_node = node.all_input_nodes[0]
-            in_quantargs = search_quant_arg_upstream(input_node)
+            in_quantargs = get_quant_arg_upstream(input_node)
             output_node = list(node.users)[0]
-            out_quantargs = search_quant_arg_downstream(output_node)
+            out_quantargs = get_quant_arg_downstream(output_node)
             table = rsqrt_table_8bit(in_quantargs, out_quantargs)
             table_attr = ts.TosaSerializerAttribute()
             table_attr.TableAttribute(table)

backends/arm/operators/op_sigmoid.py

@@ -17,10 +17,10 @@
 
 from executorch.backends.arm.tosa_quant_utils import (
     dequantize_value,
+    get_quant_arg_downstream,
+    get_quant_arg_upstream,
     QuantArgs,
     quantize_value,
-    search_quant_arg_downstream,
-    search_quant_arg_upstream,
 )
 from serializer.tosa_serializer import TosaOp
 from torch.fx import Node
@@ -50,9 +50,9 @@ def define_node(
 
             # Create attribute for 8 bit table lookup.
             input_node = node.all_input_nodes[0]
-            in_quantargs = search_quant_arg_upstream(input_node)
+            in_quantargs = get_quant_arg_upstream(input_node)
             output_node = list(node.users)[0]
-            out_quantargs = search_quant_arg_downstream(output_node)
+            out_quantargs = get_quant_arg_downstream(output_node)
 
             table = sigmoid_table_8bit(in_quantargs, out_quantargs)
             table_attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_tanh.py

@@ -17,10 +17,10 @@
 
 from executorch.backends.arm.tosa_quant_utils import (
     dequantize_value,
+    get_quant_arg_downstream,
+    get_quant_arg_upstream,
     QuantArgs,
     quantize_value,
-    search_quant_arg_downstream,
-    search_quant_arg_upstream,
 )
 from serializer.tosa_serializer import TosaOp
 from torch.fx import Node
@@ -50,9 +50,9 @@ def define_node(
 
             # Create attribute for 8 bit table lookup.
             input_node = node.all_input_nodes[0]
-            in_quantargs = search_quant_arg_upstream(input_node)
+            in_quantargs = get_quant_arg_upstream(input_node)
             output_node = list(node.users)[0]
-            out_quantargs = search_quant_arg_downstream(output_node)
+            out_quantargs = get_quant_arg_downstream(output_node)
 
             table = tanh_table_8bit(in_quantargs, out_quantargs)
             table_attr = ts.TosaSerializerAttribute()