Arm backend: Handle i48 special case for bias tensor

For the case when the activation is 16 bit the bias in TOSA must be
a int48_t tensor. Since that can't be represented using torch.dtypes
the corresponding node.meta is set with a key 'tosa_dtype_48bit' to
pass through the note to the creation of the TOSA Tensor.
Also make sure to distinguish between int32 and int48 tensors in fuse
constant ops pass.

Signed-off-by: Per Åstrand <[email protected]>
Change-Id: Iefe64f2b02f388c905c9c818ee7d2a6af40bc9e3

Author: per

backends/arm/_passes/add_bias_pass.py

@@ -55,6 +55,8 @@ def call(self, graph_module):
                         persistent_buffer=True,
                         name=f"{node.name}_bias",
                     )
+                    if node.args[0].meta["val"].dtype == torch.int16:
+                        bias_node.meta["tosa_dtype_48bit"] = True
                 node.update_arg(2, bias_node)
 
         if modified:

backends/arm/_passes/fuse_equal_placeholders_pass.py

@@ -44,9 +44,14 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                 continue
             # Create a lightweight fingerprint: dtype + shape + SHA1 of raw bytes
             # Ensure tensor is on CPU and contiguous
+
+            # ensure we don't merge any special case int48_t tensors with int32_t tensors
+            # since int48_t tensors needs to be instantiated separately.
+            is_int48 = node.meta.get("tosa_dtype_48bit", False)
             t_cpu = tensor.detach().cpu().contiguous()
             data_bytes = t_cpu.numpy().tobytes()
             key = (
+                is_int48,
                 str(t_cpu.dtype),
                 tuple(t_cpu.shape),
                 hashlib.sha1(data_bytes).hexdigest(),

backends/arm/process_node.py

@@ -119,10 +119,16 @@ def process_inputs_to_parameters(
     if tosa_arg.dtype == torch.float32:
         assert tosa_spec.support_float(), f"{tosa_spec} doesn't support float"
 
+    # Handle special case for INT48 tensors
+    if node.meta.get("tosa_dtype_48bit", False):
+        tosa_dtype = ts.DType.INT48
+    else:
+        tosa_dtype = tosa_arg.dtype
+
     parameter_values = np.transpose(parameter_values, tosa_arg.dim_order)
 
     tosa_graph.addConst(
-        parameter_values.shape, tosa_arg.dtype, parameter_values, name=tosa_arg.name
+        parameter_values.shape, tosa_dtype, parameter_values, name=tosa_arg.name
     )
 
 

backends/arm/tosa/mapping.py

@@ -81,10 +81,16 @@ def extract_tensor_meta(meta, tosa_spec: TosaSpecification):
 class TosaArg:
     def __process_node(self, argument: torch.fx.Node):
         self.name: str = argument.name
-        self.dtype, self.shape, self.dim_order = extract_tensor_meta(
+        output_dtype, self.shape, self.dim_order = extract_tensor_meta(
             argument.meta, self.tosa_spec
         )
 
+        # Handle special case of int
+        if argument.meta.get("tosa_dtype_48bit", False):
+            output_dtype = ts.DType.INT48
+
+        self.dtype = output_dtype
+
     def __process_list(self, argument):
         self.special: list = list(argument)