Arm backend: Use output.name in node visitors

As mentioned in #15381, TOSA tensors need unique naming,
which gets tricky with submodules. It is handled in
the TosaArg object, and therefore node visitors need
to use output.name rather than node.name when creating
new tensors.

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

Author: Erik-Lundell

backends/arm/operators/op_index_tensor.py

@@ -165,14 +165,14 @@ def define_node(
             # channels and thus the stride-shift.
             data = np.full(index_shape, int(values_strides[i] / C))
             mul_const = tosa_graph.addConst(index_shape, index_dtype, data)
-            tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{node.name}_{i}_shift")
+            tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{output.name}_{i}_shift")
             attr = ts.TosaSerializerAttribute()
             attr.MulAttribute()
             self._serialize_operator(
                 node,
                 tosa_graph,
                 ts.Op.MUL,
-                [index_name, mul_const.name, f"{node.name}_{i}_shift"],
+                [index_name, mul_const.name, f"{output.name}_{i}_shift"],
                 [stride_shifted_indices.name],
                 attr,
             )
@@ -186,7 +186,7 @@ def define_node(
                 stride_shifted_indices.name,
                 gather_idx_shape,
                 reshaped_idxs.name,
-                shape_name_override=f"{node.name}_{i}_shape",
+                shape_name_override=f"{output.name}_{i}_shape",
             )
 
             # Guarantees that the accumulation tensor is properly
@@ -218,7 +218,7 @@ def define_node(
             values.name,
             gather_vals_shape,
             reshaped_input.name,
-            shape_name_override=f"{node.name}_index_shape",
+            shape_name_override=f"{output.name}_index_shape",
         )
 
         gather_out_shape = (N, W, C)
@@ -244,5 +244,5 @@ def define_node(
             gather_out.name,
             list(output_shape),
             output.name,
-            shape_name_override=f"{node.name}_output_shape",
+            shape_name_override=f"{output.name}_output_shape",
         )

backends/arm/operators/op_mul.py

@@ -48,14 +48,14 @@ def define_node(
             output.tosa_spec,
         )
 
-        tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{node.name}_shift")
+        tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{output.name}_shift")
         attr = ts.TosaSerializerAttribute()
         attr.MulAttribute()
         self._serialize_operator(
             node,
             tosa_graph,
             ts.Op.MUL,
-            [inputs[0].name, inputs[1].name, f"{node.name}_shift"],
+            [inputs[0].name, inputs[1].name, f"{output.name}_shift"],
             [output.name],
             attr,
         )

backends/arm/operators/op_repeat.py

@@ -56,7 +56,7 @@ def define_node(
             (len(multiples),),
             ts.DType.SHAPE,
             list(tosa_shape(multiples, output.dim_order)),
-            name=node.name + "_multiples",
+            name=output.name + "_multiples",
         )
 
         attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_slice.py

@@ -120,7 +120,7 @@ def define_node(
             (starts_len,),
             ts.DType.SHAPE,
             starts,
-            node.name + "_start_shape",
+            output.name + "_start_shape",
         )
 
         sizes = [size if i == dim else shape[i] for i in input_node.dim_order]
@@ -130,7 +130,7 @@ def define_node(
             sizes_len = 1
             sizes = [0]
         sizes_tensor = tosa_graph.addConst(
-            (sizes_len,), ts.DType.SHAPE, sizes, node.name + "_sizes_shape"
+            (sizes_len,), ts.DType.SHAPE, sizes, output.name + "_sizes_shape"
         )
 
         attr = ts.TosaSerializerAttribute()

backends/arm/operators/op_tosa_matmul.py

@@ -72,8 +72,8 @@ def define_node(
         else:
             input0_zp, input1_zp = 0, 0
 
-        input_A_ZP_name = f"{node.name}_A_ZP"
-        input_B_ZP_name = f"{node.name}_B_ZP"
+        input_A_ZP_name = f"{output.name}_A_ZP"
+        input_B_ZP_name = f"{output.name}_B_ZP"
         tosa_graph.addConst([1], inputs[0].dtype, [input0_zp], name=input_A_ZP_name)
         tosa_graph.addConst([1], inputs[1].dtype, [input1_zp], name=input_B_ZP_name)
 

backends/arm/operators/op_tosa_resize.py

@@ -84,15 +84,15 @@ def in_int16_range(x):
             scale_d_vals[1],
         ]
         scales_tensor = tosa_graph.addConst(
-            [len(scales)], ts.DType.SHAPE, scales, node.name + "_scales"
+            [len(scales)], ts.DType.SHAPE, scales, output.name + "_scales"
         )
         offset = [int(v) for v in offset_yx.tolist()]
         offset_tensor = tosa_graph.addConst(
-            [len(offset)], ts.DType.SHAPE, offset, node.name + "_offset"
+            [len(offset)], ts.DType.SHAPE, offset, output.name + "_offset"
         )
         border = [int(v) for v in border_yx.tolist()]
         border_tensor = tosa_graph.addConst(
-            [len(border)], ts.DType.SHAPE, border, node.name + "_border"
+            [len(border)], ts.DType.SHAPE, border, output.name + "_border"
         )
         attr = ts.TosaSerializerAttribute()
         attr.ResizeAttribute(resize_mode)

backends/arm/operators/op_tosa_table.py

@@ -44,27 +44,24 @@ def define_node(
         if inputs[0].dtype == ts.DType.INT16:
             validate_valid_dtype(self.target, output, ts.DType.INT32, output.tosa_spec)
 
-        if inputs[1].name not in self._exported_program.state_dict.keys():  # type: ignore[union-attr]
+        # The name of the table constant is a bit complex.
+        # The name of the pytorch buffer will be the target of last node argument.
+        # However, when it is serialized to TOSA, a submodule suffix might be added. The TOSA buffer name thus
+        # needs to be taken from the last TosaArg.
+        pytorch_table_buffer_name = node.args[-1].target  # type: ignore[union-attr]
+        tosa_table_buffer_name = inputs[-1].name
+        if pytorch_table_buffer_name not in self._exported_program.state_dict.keys():
             raise RuntimeError(
                 f"Did not find key {node.name} in state_dict {self._exported_program.state_dict.keys()}."
             )
 
-        table = self._exported_program.state_dict[inputs[1].name]  # type: ignore[union-attr]
-
-        table_tensor_name = node.name + "_table"
-        tosa_graph.addConst(
-            table.shape,
-            ts.DType.INT8 if inputs[0].dtype == ts.DType.INT8 else ts.DType.INT16,
-            table.detach().numpy(),
-            name=table_tensor_name,
-        )
         attr = ts.TosaSerializerAttribute()
         attr.TableAttribute()
         self._serialize_operator(
             node,
             tosa_graph,
             ts.Op.TABLE,
-            [inputs[0].name, table_tensor_name],
+            [inputs[0].name, tosa_table_buffer_name],
             [output.name],
             attr,
         )

backends/arm/operators/op_view.py

@@ -66,7 +66,7 @@ def define_node(
             shape_len,
             ts.DType.SHAPE,
             shape_data,
-            name=node.name + "_shape",
+            name=output.name + "_shape",
         )
 
         attr = ts.TosaSerializerAttribute()