Arm backend: Add special dtype TOSA handling

Add a enum class to handle special dtypes that can't be represented
in torch (i.e. int48_t) to avoid leaking serializer types into the
pass handling of the backend.

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

Author: per

backends/arm/_passes/add_bias_pass.py

@@ -8,6 +8,7 @@
 import torch
 from executorch.backends.arm._passes import ArmPass
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
+from executorch.backends.arm.tosa.mapping import TosaSpecialDtype
 from executorch.backends.transforms.utils import create_constant_placeholder
 
 from executorch.exir.dialects._ops import ops as exir_ops
@@ -60,7 +61,9 @@ def call(self, graph_module):
                         name=f"{node.name}_bias",
                     )
                     if node.args[0].meta["val"].dtype == torch.int16:
-                        bias_node.meta["tosa_dtype_48bit"] = True
+                        bias_node.meta[TosaSpecialDtype.meta_key()] = (
+                            TosaSpecialDtype.INT48
+                        )
                 node.update_arg(2, bias_node)
 
         if modified:

backends/arm/_passes/fuse_equal_placeholders_pass.py

@@ -8,11 +8,13 @@
 from typing import Set, Type
 
 import torch
+
 from executorch.backends.arm._passes.arm_pass_utils import (
     get_constant_placeholder_kind,
     get_param_tensor,
     is_param_node,
 )
+from executorch.backends.arm.tosa.mapping import TosaSpecialDtype
 from executorch.backends.transforms.utils import (
     create_constant_placeholder,
     delete_constant_placeholder,
@@ -50,7 +52,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 
             # 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)
+            is_int48 = node.meta.get(TosaSpecialDtype.meta_key(), None)
             t_cpu = tensor.detach().cpu().contiguous()
             data_bytes = t_cpu.numpy().tobytes()
             key = (

backends/arm/process_node.py

@@ -12,7 +12,7 @@
 import torch
 import torch.fx
 from executorch.backends.arm.operators.node_visitor import NodeVisitor
-from executorch.backends.arm.tosa.mapping import TosaArg
+from executorch.backends.arm.tosa.mapping import TosaArg, TosaSpecialDtype
 from executorch.backends.arm.tosa.specification import TosaSpecification
 from executorch.backends.arm.tosa.utils import tosa_shape
 from torch._export.utils import (
@@ -113,8 +113,9 @@ def process_inputs_to_parameters(
         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
+    special_type = node.meta.get(TosaSpecialDtype.meta_key(), None)
+    if isinstance(special_type, TosaSpecialDtype):
+        tosa_dtype = special_type.get_tosa_dtype()
     else:
         tosa_dtype = tosa_arg.dtype
 

backends/arm/tosa/mapping.py

@@ -11,6 +11,7 @@
 
 """
 
+from enum import Enum
 from typing import Any, Optional, Sequence
 
 import serializer.tosa_serializer as ts  # type: ignore
@@ -31,6 +32,22 @@
 )
 
 
+class TosaSpecialDtype(Enum):
+    """
+    Special TOSA data types that are not natively supported in PyTorch, to be
+    used in specific scenarios as a value in the key from meta_key().
+    """
+
+    INT48 = ts.DType.INT48
+
+    def get_tosa_dtype(self) -> ts.TosaDType.DType:
+        return self.value
+
+    @staticmethod
+    def meta_key() -> str:
+        return "tosa_special_dtype"
+
+
 def map_dtype(data_type: torch.dtype, tosa_spec: TosaSpecification) -> Any:
     """Map a ``torch.dtype`` to a ``ts.DType``.
 
@@ -134,9 +151,9 @@ def __process_node(self, argument: torch.fx.Node):
             argument.meta, self.tosa_spec
         )
 
-        # Handle special case of int
-        if argument.meta.get("tosa_dtype_48bit", False):
-            output_dtype = ts.DType.INT48
+        # Handle special case of types not representable in torch (i.e. i48_t)
+        if special_type := argument.meta.get(TosaSpecialDtype.meta_key(), None):
+            output_dtype = special_type.get_tosa_dtype()
 
         self.dtype = output_dtype