Arm backend: Fix TOSA 1.0 node visitor for sum (#10908)

per · web-flow · commit 12af535a5e5b · 2025-05-16T11:33:44.000+02:00
### Summary
Fixes serialization for sum.dim_IntList node visitor as well as some
rescale handling issues.

### Test plan
Tested with internal and external GitHub CI.

Signed-off-by: Per Åstrand &lt;per.astrand@arm.com&gt;
diff --git a/backends/arm/operators/op_abs.py b/backends/arm/operators/op_abs.py
@@ -164,7 +164,7 @@ def define_node(
         scale_back = 1.0
         if inputs[0].dtype == ts.DType.INT8:
             rescaled_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(
-                tosa_graph, inputs, node, self.tosa_specs
+                tosa_graph, inputs, node, self.tosa_spec
             )  # type: ignore[possibly-undefined]
         else:
             # input[0].dtype == ts.DType.INT32
@@ -192,7 +192,7 @@ def define_node(
             # Scale output back to 8 bit
             # pyre-ignore
             tqutils.insert_rescale_op_to_int8(
-                tosa_graph, abs_output, scale_back, node, self.tosa_specs
+                tosa_graph, abs_output, scale_back, node, self.tosa_spec
             )  # type: ignore[possibly-undefined]
 
 
diff --git a/backends/arm/operators/op_add.py b/backends/arm/operators/op_add.py
@@ -174,7 +174,7 @@ def define_node(
         scale_back = 1.0
         if inputs[0].dtype == ts.DType.INT8:
             rescaled_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(
-                tosa_graph, inputs, node, self.tosa_specs
+                tosa_graph, inputs, node, self.tosa_spec
             )
         else:
             # input[0].dtype == ts.DType.INT32
@@ -202,7 +202,7 @@ def define_node(
             # Scale output back to 8 bit
             # pyre-ignore
             tqutils.insert_rescale_op_to_int8(
-                tosa_graph, add_output, scale_back, node, self.tosa_specs
+                tosa_graph, add_output, scale_back, node, self.tosa_spec
             )  # type: ignore[possibly-undefined]
 
 
diff --git a/backends/arm/operators/op_eq.py b/backends/arm/operators/op_eq.py
@@ -98,7 +98,7 @@ def define_node(
         if inputs[0].dtype == ts.DType.INT8:
             # Rescale inputs to 32 bit
             rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
-                tosa_graph, inputs, node, self.tosa_specs
+                tosa_graph, inputs, node, self.tosa_spec
             )
 
             # Update IO
diff --git a/backends/arm/operators/op_ge.py b/backends/arm/operators/op_ge.py
@@ -97,7 +97,7 @@ def define_node(
         if inputs[0].dtype == ts.DType.INT8:
             # Rescale inputs to 32 bit
             rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
-                tosa_graph, inputs, node, self.tosa_specs
+                tosa_graph, inputs, node, self.tosa_spec
             )
 
             # Update IO
diff --git a/backends/arm/operators/op_gt.py b/backends/arm/operators/op_gt.py
@@ -97,7 +97,7 @@ def define_node(
         if inputs[0].dtype == ts.DType.INT8:
             # Rescale inputs to 32 bit
             rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
-                tosa_graph, inputs, node, self.tosa_specs
+                tosa_graph, inputs, node, self.tosa_spec
             )
 
             # Update IO
diff --git a/backends/arm/operators/op_le.py b/backends/arm/operators/op_le.py
@@ -97,7 +97,7 @@ def define_node(
         if inputs[0].dtype == ts.DType.INT8:
             # Rescale inputs to 32 bit
             rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
-                tosa_graph, inputs, node, self.tosa_specs
+                tosa_graph, inputs, node, self.tosa_spec
             )
 
             # Update IO
diff --git a/backends/arm/operators/op_lt.py b/backends/arm/operators/op_lt.py
@@ -97,7 +97,7 @@ def define_node(
         if inputs[0].dtype == ts.DType.INT8:
             # Rescale inputs to 32 bit
             rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(
-                tosa_graph, inputs, node, self.tosa_specs
+                tosa_graph, inputs, node, self.tosa_spec
             )
 
             # Update IO
diff --git a/backends/arm/operators/op_maximum.py b/backends/arm/operators/op_maximum.py
@@ -129,7 +129,7 @@ def define_node(
                 )
 
             operand_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(
-                tosa_graph, inputs, node, self.tosa_specs
+                tosa_graph, inputs, node, self.tosa_spec
             )
 
             output.shape = tosa_shape(output.shape, output.dim_order)
@@ -155,5 +155,5 @@ def define_node(
         if output.dtype == ts.DType.INT8:
             # insert RESCALE from int32 back to int8
             tqutils.insert_rescale_op_to_int8(
-                tosa_graph, max_output, scale_back, node, self.tosa_specs
+                tosa_graph, max_output, scale_back, node, self.tosa_spec
             )
diff --git a/backends/arm/operators/op_minimum.py b/backends/arm/operators/op_minimum.py
@@ -128,7 +128,7 @@ def define_node(
                 )
 
             operand_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(
-                tosa_graph, inputs, node, self.tosa_specs
+                tosa_graph, inputs, node, self.tosa_spec
             )
 
             output.shape = tosa_shape(output.shape, output.dim_order)
@@ -154,5 +154,5 @@ def define_node(
         if output.dtype == ts.DType.INT8:
             # insert RESCALE from int32 back to int8
             tqutils.insert_rescale_op_to_int8(
-                tosa_graph, min_output, scale_back, node, self.tosa_specs
+                tosa_graph, min_output, scale_back, node, self.tosa_spec
             )
diff --git a/backends/arm/operators/op_mul.py b/backends/arm/operators/op_mul.py
@@ -189,14 +189,14 @@ def define_node(
             input_A,
             input_A_qargs.zp,
             [1.0],
-            tosa_spec=self.tosa_specs,
+            tosa_spec=self.tosa_spec,
         )
         input_B_rescaled = tqutils.build_rescale_to_int32(
             tosa_graph,
             input_B,
             input_B_qargs.zp,
             [1.0],
-            tosa_spec=self.tosa_specs,
+            tosa_spec=self.tosa_spec,
         )
 
         output_shape = tutils.tosa_shape(output.shape, output.dim_order)
@@ -211,7 +211,7 @@ def define_node(
         )
         output_scale = input_A_qargs.scale * input_B_qargs.scale
         tqutils.insert_rescale_op_to_int8(
-            tosa_graph, mul_output, output_scale, node, self.tosa_specs
+            tosa_graph, mul_output, output_scale, node, self.tosa_spec
         )
 
 
diff --git a/backends/arm/operators/op_sub.py b/backends/arm/operators/op_sub.py
@@ -168,7 +168,7 @@ def define_node(
         scale_back = 1.0
         if inputs[0].dtype == ts.DType.INT8:
             rescaled_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(
-                tosa_graph, inputs, node, self.tosa_specs
+                tosa_graph, inputs, node, self.tosa_spec
             )
         else:
             # input[0].dtype == ts.DType.INT32
@@ -197,7 +197,7 @@ def define_node(
             # Scale output back to 8 bit
             # pyre-ignore
             tqutils.insert_rescale_op_to_int8(
-                tosa_graph, sub_output, scale_back, node, self.tosa_specs
+                tosa_graph, sub_output, scale_back, node, self.tosa_spec
             )  # type: ignore[possibly-undefined]
 
 
diff --git a/backends/arm/operators/op_sum.py b/backends/arm/operators/op_sum.py
@@ -159,13 +159,11 @@ def define_node(
 
         # Rescale input to 32 bit
         rescaled_inputs, scale = tqutils.insert_rescale_ops_to_int32(
-            tosa_graph,
-            [tensor],
-            node,
+            tosa_graph, [tensor], node, self.tosa_spec
         )
 
         attr = ts.TosaSerializerAttribute()
-        attr.AxisAttribute(tensor.dim_order.index(dim))
+        attr.ReduceSumAttribute(tensor.dim_order.index(dim))
 
         intermediate = tosa_graph.addIntermediate(
             tutils.tosa_shape(output_shape, tensor.dim_order),
@@ -179,7 +177,9 @@ def define_node(
             attr,
         )
 
-        tqutils.insert_rescale_op_to_int8(tosa_graph, intermediate, scale, node)
+        tqutils.insert_rescale_op_to_int8(
+            tosa_graph, intermediate, scale, node, self.tosa_spec
+        )
 
 
 @register_node_visitor
@@ -212,7 +212,7 @@ def define_node(
         output_shape[dim] = 1  # Output shape is input shape with dim reduced
 
         attr = ts.TosaSerializerAttribute()
-        attr.AxisAttribute(tensor.dim_order.index(dim))
+        attr.ReduceSumAttribute(tensor.dim_order.index(dim))
 
         tosa_graph.addOperator(
             ts.TosaOp.Op().REDUCE_SUM,
diff --git a/backends/arm/tosa_quant_utils.py b/backends/arm/tosa_quant_utils.py
@@ -32,7 +32,7 @@ def insert_rescale_ops_to_int32(
     tosa_graph: Any,
     inputs: list[TosaArg],
     node: Node,
-    tosa_spec=tosa_specification.Tosa_0_80,
+    tosa_spec=None,
 ) -> tuple[list[Any], float]:
     """Rescales all 'nodes' to int32, adding suitable RESCALE ops to 'tosa_graph'.
     The scales are adjusted using the smallest scale of all 'nodes'.
@@ -79,7 +79,7 @@ def insert_rescale_op_to_int8(
     last_tensor: TosaArg,
     scale: float,
     node: Node,
-    tosa_spec=tosa_specification.Tosa_0_80,
+    tosa_spec=None,
 ) -> None:
     """Rescales the node back to int8, adding a suitable RESCALE op to 'tosa_graph'.
     Parameters:
@@ -323,10 +323,11 @@ def build_rescale_to_int32(
     is_scale32: bool = True,
     is_double_round: bool = False,
     per_channel: bool = False,
-    tosa_spec=tosa_specification.Tosa_0_80,
+    tosa_spec=None,
 ) -> Any:
     input_A_rescaled_to_int32 = None
-    if tosa_spec == tosa_specification.Tosa_0_80:
+    if not tosa_spec or isinstance(tosa_spec, tosa_specification.Tosa_0_80):
+        # default to TOSA v0.80 until we switch to v1.0
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
         input_A_rescaled_to_int32 = tosa_fb.addIntermediate(
@@ -343,7 +344,7 @@ def build_rescale_to_int32(
             output_zp=0,
         )  # type: ignore[call-arg]
 
-    elif isinstance(tosa_spec[0], tosa_specification.Tosa_1_00):
+    elif isinstance(tosa_spec, tosa_specification.Tosa_1_00):
         # For TOSA v1.0 multipliers, shifts, input_zp and output_zp are now inputs
         # to the RESCALE op see: https://www.mlplatform.org/tosa/tosa_spec.html#_rescale
         import serializer.tosa_serializer as ts  # type: ignore
@@ -375,9 +376,10 @@ def build_rescale_from_int32(
     is_scale32: bool = True,
     is_double_round: bool = False,
     per_channel: bool = False,
-    tosa_spec=tosa_specification.Tosa_0_80,
+    tosa_spec=None,
 ) -> None:
-    if tosa_spec == tosa_specification.Tosa_0_80:
+    if not tosa_spec or isinstance(tosa_spec, tosa_specification.Tosa_0_80):
+        # default to TOSA v0.80 until we switch to v1.0
         import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
         build_rescale_v0_80(
@@ -390,7 +392,7 @@ def build_rescale_from_int32(
             output_zp=output_zp,
         )  # type: ignore[call-arg]
 
-    elif isinstance(tosa_spec[0], tosa_specification.Tosa_1_00):
+    elif isinstance(tosa_spec, tosa_specification.Tosa_1_00):
         import serializer.tosa_serializer as ts  # type: ignore
 
         # For TOSA v1.0 multipliers, shifts, input_zp and output_zp are now inputs
@@ -420,15 +422,16 @@ def build_rescale_conv_output(
     weight_scale: list[float],
     output_scale: list[float],
     output_zp: int,
-    tosa_spec=tosa_specification.Tosa_0_80,
+    tosa_spec=None,
 ):
     # TODO add check to verify if this is a Per-channel quantization.
     post_conv2d_scale = [
         (inp * w) / out for inp, w, out in zip(input_scale, weight_scale, output_scale)
     ]
 
     # Since we assume the input tensor that is being rescaled is int32 date type, zero point must be 0.
-    if tosa_spec == tosa_specification.Tosa_0_80:
+    if not tosa_spec or isinstance(tosa_spec, tosa_specification.Tosa_0_80):
+        # default to TOSA v0.80 until we switch to v1.0
         build_rescale_v0_80(
             tosa_fb=tosa_fb,
             scale=post_conv2d_scale,

Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,7 @@ def define_node(`
`98`	`98`	`if inputs[0].dtype == ts.DType.INT8:`
`99`	`99`	`# Rescale inputs to 32 bit`
`100`	`100`	`rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(`
`101`		`- tosa_graph, inputs, node, self.tosa_specs`
	`101`	`+ tosa_graph, inputs, node, self.tosa_spec`
`102`	`102`	`)`
`103`	`103`
`104`	`104`	`# Update IO`
Original file line number	Diff line number	Diff line change
`@@ -97,7 +97,7 @@ def define_node(`
`97`	`97`	`if inputs[0].dtype == ts.DType.INT8:`
`98`	`98`	`# Rescale inputs to 32 bit`
`99`	`99`	`rescaled_inputs, _ = tqutils.insert_rescale_ops_to_int32(`
`100`		`- tosa_graph, inputs, node, self.tosa_specs`
	`100`	`+ tosa_graph, inputs, node, self.tosa_spec`
`101`	`101`	`)`
`102`	`102`
`103`	`103`	`# Update IO`
Original file line number	Diff line number	Diff line change
`@@ -129,7 +129,7 @@ def define_node(`
`129`	`129`	`)`
`130`	`130`
`131`	`131`	`operand_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(`
`132`		`- tosa_graph, inputs, node, self.tosa_specs`
	`132`	`+ tosa_graph, inputs, node, self.tosa_spec`
`133`	`133`	`)`
`134`	`134`
`135`	`135`	`output.shape = tosa_shape(output.shape, output.dim_order)`
`@@ -155,5 +155,5 @@ def define_node(`
`155`	`155`	`if output.dtype == ts.DType.INT8:`
`156`	`156`	`# insert RESCALE from int32 back to int8`
`157`	`157`	`tqutils.insert_rescale_op_to_int8(`
`158`		`- tosa_graph, max_output, scale_back, node, self.tosa_specs`
	`158`	`+ tosa_graph, max_output, scale_back, node, self.tosa_spec`
`159`	`159`	`)`
Original file line number	Diff line number	Diff line change
`@@ -128,7 +128,7 @@ def define_node(`
`128`	`128`	`)`
`129`	`129`
`130`	`130`	`operand_inputs, scale_back = tqutils.insert_rescale_ops_to_int32(`
`131`		`- tosa_graph, inputs, node, self.tosa_specs`
	`131`	`+ tosa_graph, inputs, node, self.tosa_spec`
`132`	`132`	`)`
`133`	`133`
`134`	`134`	`output.shape = tosa_shape(output.shape, output.dim_order)`
`@@ -154,5 +154,5 @@ def define_node(`
`154`	`154`	`if output.dtype == ts.DType.INT8:`
`155`	`155`	`# insert RESCALE from int32 back to int8`
`156`	`156`	`tqutils.insert_rescale_op_to_int8(`
`157`		`- tosa_graph, min_output, scale_back, node, self.tosa_specs`
	`157`	`+ tosa_graph, min_output, scale_back, node, self.tosa_spec`
`158`	`158`	`)`
Original file line number	Diff line number	Diff line change
`@@ -189,14 +189,14 @@ def define_node(`
`189`	`189`	`input_A,`
`190`	`190`	`input_A_qargs.zp,`
`191`	`191`	`[1.0],`
`192`		`- tosa_spec=self.tosa_specs,`
	`192`	`+ tosa_spec=self.tosa_spec,`
`193`	`193`	`)`
`194`	`194`	`input_B_rescaled = tqutils.build_rescale_to_int32(`
`195`	`195`	`tosa_graph,`
`196`	`196`	`input_B,`
`197`	`197`	`input_B_qargs.zp,`
`198`	`198`	`[1.0],`
`199`		`- tosa_spec=self.tosa_specs,`
	`199`	`+ tosa_spec=self.tosa_spec,`
`200`	`200`	`)`
`201`	`201`
`202`	`202`	`output_shape = tutils.tosa_shape(output.shape, output.dim_order)`
`@@ -211,7 +211,7 @@ def define_node(`
`211`	`211`	`)`
`212`	`212`	`output_scale = input_A_qargs.scale * input_B_qargs.scale`
`213`	`213`	`tqutils.insert_rescale_op_to_int8(`
`214`		`- tosa_graph, mul_output, output_scale, node, self.tosa_specs`
	`214`	`+ tosa_graph, mul_output, output_scale, node, self.tosa_spec`
`215`	`215`	`)`
`216`	`216`
`217`	`217`