Fix type-checking issues in Arm backend

Erik-Lundell · Erik-Lundell · commit 8455a71349f4 · 2024-11-11T12:47:33.000+01:00
Signed-off-by: Erik Lundell &lt;erik.lundell@arm.com&gt;
Change-Id: I4feb4b5c6d269d7c0ff4312c17fec52da413fa5a
diff --git a/backends/arm/operators/op_bmm.py b/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
diff --git a/backends/arm/operators/op_conv2d.py b/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,
diff --git a/backends/arm/operators/op_exp.py b/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()
diff --git a/backends/arm/operators/op_full.py b/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)
diff --git a/backends/arm/operators/op_hardtanh.py b/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)
diff --git a/backends/arm/operators/op_log.py b/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()
diff --git a/backends/arm/operators/op_max_pool2d.py b/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(
diff --git a/backends/arm/operators/op_mm.py b/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
diff --git a/backends/arm/operators/op_mul.py b/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])
             )
 
diff --git a/backends/arm/operators/op_placeholder.py b/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 (
@@ -68,8 +68,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()
diff --git a/backends/arm/operators/op_reciprocal.py b/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)
 
diff --git a/backends/arm/operators/op_relu.py b/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)
 
diff --git a/backends/arm/operators/op_rsqrt.py b/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)
diff --git a/backends/arm/operators/op_sigmoid.py b/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()
diff --git a/backends/arm/operators/op_tanh.py b/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()
diff --git a/backends/arm/quantizer/arm_quantizer.py b/backends/arm/quantizer/arm_quantizer.py
@@ -75,7 +75,7 @@ def _supported_symmetric_quantized_operators() -> Dict[str, List[OperatorPattern
             [torch.nn.AdaptiveAvgPool2d],
             [F.adaptive_avg_pool2d],
         ],
-        "mul": [torch.mul],
+        "mul": [[torch.mul]],
         "sub": [[torch.sub]],
     }
     return copy.deepcopy(supported_operators)
diff --git a/backends/arm/tosa_quant_utils.py b/backends/arm/tosa_quant_utils.py
diff --git a/backends/arm/tosa_utils.py b/backends/arm/tosa_utils.py
diff --git a/backends/arm/util/arm_model_evaluator.py b/backends/arm/util/arm_model_evaluator.py

Original file line number	Diff line number	Diff line change
`@@ -37,10 +37,10 @@ def define_node(`
`37`	`37`	`if is_quant_node:`
`38`	`38`	`input_A = inputs[0]`
`39`	`39`	`input_B = inputs[1]`
`40`		`- input_A_qargs = tqutils.search_quant_arg_upstream(`
	`40`	`+ input_A_qargs = tqutils.get_quant_arg_upstream(`
`41`	`41`	`cast(torch.fx.Node, node.args[0])`
`42`	`42`	`)`
`43`		`- input_B_qargs = tqutils.search_quant_arg_upstream(`
	`43`	`+ input_B_qargs = tqutils.get_quant_arg_upstream(`
`44`	`44`	`cast(torch.fx.Node, node.args[1])`
`45`	`45`	`)`
`46`	`46`