Update

[ghstack-poisoned]

Author: JulianCloudNTH

backends/arm/_passes/rewrite_avg_pool2d_pass.py

@@ -65,9 +65,11 @@ def call_operator(self, op, args, kwargs, meta, updated=False):
         # Materialize output zero-point as a scalar tensor
         output_zp = super().call_scalar(out_zp_val, meta)
 
-        # Determine accumulator dtype for AVG_POOL2D: INT32 for integer inputs, FP32 otherwise
+        # Determine accumulator dtype for AVG_POOL2D.
         if x.data.dtype in (torch.int8, torch.int16):
             acc_type = torch.int32
+        elif x.data.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+            acc_type = torch.float16
         else:
             acc_type = torch.float32
 

backends/arm/_passes/rewrite_conv_pass.py

@@ -5,7 +5,7 @@
 
 
 import itertools
-from typing import Any, Set, Type
+from typing import Any, cast, Set, Type
 
 import torch
 from executorch.backends.arm._passes import ArmPass
@@ -39,6 +39,7 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 
+from torch._subclasses.fake_tensor import FakeTensor
 from torch.export.graph_signature import InputKind
 
 
@@ -350,6 +351,68 @@ def _has_int32_rescale_user(self, node: torch.fx.Node) -> bool:
                         return True
         return False
 
+    def _insert_output_conversion(
+        self,
+        graph_module: torch.fx.GraphModule,
+        node: torch.fx.Node,
+        tosa_op: torch.fx.Node,
+        input_fake_tensor: torch.Tensor,
+        tosa_node_fake_tensor: torch.Tensor,
+    ) -> tuple[torch.fx.Node, FakeTensor]:
+        node_replacement: torch.fx.Node = tosa_op
+        node_replacement_fake_tensor = tosa_node_fake_tensor
+        if (
+            tosa_node_fake_tensor.dtype == torch.int32
+            and input_fake_tensor.dtype == torch.int8
+        ):
+            node_replacement, node_replacement_fake_tensor = self.insert_output_rescale(
+                graph_module, node, tosa_op, tosa_node_fake_tensor
+            )
+        elif (
+            tosa_node_fake_tensor.dtype == torch.int32
+            and input_fake_tensor.dtype == torch.int16
+        ):
+            # Explicit layout paths require a post-conv permute, which does
+            # not support INT48. Always rescale before post-permute.
+            if self._has_int32_rescale_user(node):
+                node_replacement, node_replacement_fake_tensor = (
+                    self.insert_identity_int32_rescale(
+                        graph_module, node, tosa_op, tosa_node_fake_tensor
+                    )
+                )
+            else:
+                node_replacement, node_replacement_fake_tensor = (
+                    self.insert_output_rescale(
+                        graph_module, node, tosa_op, tosa_node_fake_tensor
+                    )
+                )
+
+            tosa_op.meta[TosaSpecialDtype.meta_key()] = TosaSpecialDtype.INT48
+        elif (
+            tosa_node_fake_tensor.dtype == torch.float16
+            and input_fake_tensor.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
+        ):
+            node_output_fake_tensor = get_first_fake_tensor(node)
+            # TOSA FP8 conv widens the output. Cast back to the exported
+            # graph dtype before the post-layout permute.
+            node_replacement_fake_tensor = (
+                exir_ops.edge.dim_order_ops._to_dim_order_copy.default(
+                    tosa_node_fake_tensor,
+                    dtype=node_output_fake_tensor.dtype,
+                )
+            )
+            with graph_module.graph.inserting_after(tosa_op):
+                node_replacement = create_node(
+                    graph=graph_module.graph,
+                    op_target=exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
+                    args=(tosa_op,),
+                    kwargs={"dtype": node_output_fake_tensor.dtype},
+                    from_node=tosa_op,
+                )
+            node_replacement.meta["val"] = node_replacement_fake_tensor
+
+        return node_replacement, cast(FakeTensor, node_replacement_fake_tensor)
+
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
         modified = False
         for node in graph_module.graph.nodes:
@@ -561,37 +624,15 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
             )
             tosa_op.meta["val"] = tosa_node_fake_tensor
 
-            node_replacement: torch.fx.Node = tosa_op
-            node_replacement_fake_tensor = tosa_node_fake_tensor
-            if (
-                tosa_node_fake_tensor.dtype == torch.int32
-                and input_fake_tensor.dtype == torch.int8
-            ):
-                output_rescale, output_rescale_fake = self.insert_output_rescale(
-                    graph_module, node, tosa_op, tosa_node_fake_tensor
+            node_replacement, node_replacement_fake_tensor = (
+                self._insert_output_conversion(
+                    graph_module,
+                    node,
+                    tosa_op,
+                    input_fake_tensor,
+                    tosa_node_fake_tensor,
                 )
-                node_replacement = output_rescale
-                node_replacement_fake_tensor = output_rescale_fake
-            elif (
-                tosa_node_fake_tensor.dtype == torch.int32
-                and input_fake_tensor.dtype == torch.int16
-            ):
-                # Explicit layout paths require a post-conv permute, which does
-                # not support INT48. Always rescale before post-permute.
-                if self._has_int32_rescale_user(node):
-                    output_rescale, output_rescale_fake = (
-                        self.insert_identity_int32_rescale(
-                            graph_module, node, tosa_op, tosa_node_fake_tensor
-                        )
-                    )
-                else:
-                    output_rescale, output_rescale_fake = self.insert_output_rescale(
-                        graph_module, node, tosa_op, tosa_node_fake_tensor
-                    )
-                node_replacement = output_rescale
-                node_replacement_fake_tensor = output_rescale_fake
-
-                tosa_op.meta[TosaSpecialDtype.meta_key()] = TosaSpecialDtype.INT48
+            )
 
             if post_permute_dims is None:
                 raise RuntimeError("Expected post permute dims for explicit layout")

backends/arm/_passes/rewrite_matmul.py

@@ -21,12 +21,21 @@
 
 
 class RewriteMatmulPass(ArmPass):
-    """Rewrites aten.bmm to tosa.MATMUL and inserts a tosa.RESCALE op if
+    """Rewrites aten.bmm to tosa.MATMUL and inserts a tosa.RESCALE or cast op if
     needed.
     """
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
+    # TOSA MATMUL widens these floating-point input types, so outputs may need
+    # casting back to preserve the original PyTorch node semantics.
+    _WIDENING_INPUT_DTYPES = (
+        torch.float16,
+        torch.bfloat16,
+        torch.float8_e4m3fn,
+        torch.float8_e5m2,
+    )
+
     def _insert_output_rescale(self, graph_module, node, tosa_matmul_node, dtype):
         input_qparams = get_input_qparams(node)
         output_qparams = get_output_qparams(node)[0]
@@ -94,17 +103,18 @@ def call(self, graph_module):
                         TosaSpecialDtype.INT48
                     )
             elif (
-                x1_fake_tensor.dtype in [torch.float16, torch.bfloat16]
-                and x2_fake_tensor.dtype in [torch.float16, torch.bfloat16]
-                and output_fake_tensor.dtype not in [torch.float16, torch.bfloat16]
+                x1_fake_tensor.dtype in self._WIDENING_INPUT_DTYPES
+                and x2_fake_tensor.dtype in self._WIDENING_INPUT_DTYPES
+                and output_fake_tensor.dtype not in self._WIDENING_INPUT_DTYPES
             ):
-                # A TOSA BF16/FP16 MATMUL outputs FP32 whereas pytorch outputs BF16/FP16.
-                # Cast back to BF16/FP16 to get matching semantics.
+                # TOSA BF16/FP16/FP8 MATMUL outputs FP32, while the original
+                # exported node outputs BF16/FP16/FP8. Cast back to preserve
+                # the exported graph dtype.
                 with graph_module.graph.inserting_after(tosa_matmul_node):
                     cast_node = create_node(
                         graph_module.graph,
                         op_target=exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
-                        kwargs={"dtype": x1_fake_tensor.dtype},
+                        kwargs={"dtype": node_output_fake_tensor.dtype},
                         from_node=tosa_matmul_node,
                     )
                     tosa_matmul_node.replace_all_uses_with(cast_node)

backends/arm/operators/op_tosa_avg_pool2d.py

@@ -43,6 +43,10 @@ def define_node(
 
         if self.tosa_spec.support_extension("int16"):
             supported.append(ts.DType.INT16)
+        if self.tosa_spec.support_extension("fp8e4m3"):
+            supported.append(ts.DType.FP8E4M3)
+        if self.tosa_spec.support_extension("fp8e5m2"):
+            supported.append(ts.DType.FP8E5M2)
 
         validate_valid_dtype(self.target, [input, output], supported, self.tosa_spec)
 

backends/arm/operators/op_tosa_conv2d.py

@@ -67,6 +67,10 @@ def define_node(
                 )
         if self.tosa_spec.support_extension("bf16"):
             valid_input_dtypes.append(ts.DType.BF16)
+        if self.tosa_spec.support_extension("fp8e4m3"):
+            valid_input_dtypes.append(ts.DType.FP8E4M3)
+        if self.tosa_spec.support_extension("fp8e5m2"):
+            valid_input_dtypes.append(ts.DType.FP8E5M2)
 
         validate_valid_dtype(
             self.target,
@@ -82,8 +86,13 @@ def define_node(
 
         conv2d_output_name = output.name
         acc_type = output.dtype
-        if output.dtype in [ts.DType.BF16, ts.DType.FP16]:
-            # Accumulate BF16, FP16 inputs in FP32 for better precision.
+        if input.dtype in [ts.DType.FP8E4M3, ts.DType.FP8E5M2]:
+            acc_type = ts.DType.FP16
+        elif output.dtype in [
+            ts.DType.BF16,
+            ts.DType.FP16,
+        ]:
+            # Accumulate BF16 and FP16 inputs in FP32 for better precision.
             acc_type = ts.DType.FP32
 
         input_zp_name, weight_zp_name = add_input_weight_zp_consts(

backends/arm/operators/op_tosa_matmul.py

@@ -54,6 +54,8 @@ def define_node(
                 ts.DType.FP16,
                 ts.DType.FP32,
                 ts.DType.BF16,
+                ts.DType.FP8E4M3,
+                ts.DType.FP8E5M2,
             ],
             self.tosa_spec,
         )

backends/arm/operators/op_tosa_max_pool2d.py

@@ -42,6 +42,10 @@ def define_node(
         supported_dtypes = [ts.DType.INT8, ts.DType.FP16, ts.DType.FP32, ts.DType.BF16]
         if self.tosa_spec.support_extension("int16"):
             supported_dtypes.append(ts.DType.INT16)
+        if self.tosa_spec.support_extension("fp8e4m3"):
+            supported_dtypes.append(ts.DType.FP8E4M3)
+        if self.tosa_spec.support_extension("fp8e5m2"):
+            supported_dtypes.append(ts.DType.FP8E5M2)
         validate_valid_dtype(
             self.target,
             [input_tensor, output],

backends/arm/operators/op_tosa_transpose_conv2d.py

@@ -73,6 +73,24 @@ def define_node(
                 validate_valid_dtype(
                     self.target, [inputs[2]], [ts.DType.BF16], self.tosa_spec
                 )
+        if self.tosa_spec.support_extension("fp8e4m3"):
+            valid_input_dtypes.append(ts.DType.FP8E4M3)
+            if inputs[0].dtype == ts.DType.FP8E4M3:
+                validate_valid_dtype(
+                    self.target, [inputs[1]], [ts.DType.FP8E4M3], self.tosa_spec
+                )
+                validate_valid_dtype(
+                    self.target, [inputs[2]], [ts.DType.FP8E4M3], self.tosa_spec
+                )
+        if self.tosa_spec.support_extension("fp8e5m2"):
+            valid_input_dtypes.append(ts.DType.FP8E5M2)
+            if inputs[0].dtype == ts.DType.FP8E5M2:
+                validate_valid_dtype(
+                    self.target, [inputs[1]], [ts.DType.FP8E5M2], self.tosa_spec
+                )
+                validate_valid_dtype(
+                    self.target, [inputs[2]], [ts.DType.FP8E5M2], self.tosa_spec
+                )
 
         validate_valid_dtype(
             self.target,

backends/arm/scripts/aot_arm_compiler.py

@@ -1034,6 +1034,7 @@ def main() -> None:  # noqa: C901
         args.calibration_data, example_inputs
     )
     model = original_model.eval()
+    model.requires_grad_(False)
 
     # export under the assumption we quantize, the exported form also works
     # in to_edge if we don't quantize
@@ -1115,8 +1116,6 @@ def main() -> None:  # noqa: C901
 
     dump_delegation_info(edge, args.intermediates)
 
-    edge_program_manager_copy = copy.deepcopy(edge)
-
     try:
         exec_prog = edge.to_executorch(
             config=ExecutorchBackendConfig(extract_delegate_segments=False)
@@ -1175,6 +1174,7 @@ def main() -> None:  # noqa: C901
     if args.bundleio or args.etrecord:
         etrecord_file_name = os.path.splitext(output_file_name)[0] + "_etrecord.bin"
         try:
+            edge_program_manager_copy = copy.deepcopy(edge)
             generate_etrecord(etrecord_file_name, edge_program_manager_copy, exec_prog)
             print(f"ETRecord saved as {etrecord_file_name}")
         except Exception as e:

backends/arm/scripts/build_executor_runner.sh

@@ -54,7 +54,7 @@ help() {
     echo "  --et_build_root=<FOLDER>             Build output root folder to use, defaults to ${et_build_root}"
     echo "  --ethosu_tools_dir=<FOLDER>          Path to your Ethos-U tools dir if you not using default: ${ethosu_tools_dir}"
     echo "  --toolchain=<TOOLCHAIN>              Toolchain can be specified (arm-none-eabi-gcc, arm-zephyr-eabi-gcc). Default: ${toolchain}"
-    echo "  --select_ops_list=<OPS>              Comma separated list of portable (non delagated) kernels to include Default: ${select_ops_list}"
+    echo "  --select_ops_list=<OPS>              Comma separated list of portable (non-delegated) kernels to include Default: ${select_ops_list}"
     echo "                                         NOTE: This is used when select_ops_model is not possible to use, e.g. for semihosting or bundleio."
     echo "                                         See https://docs.pytorch.org/executorch/stable/kernel-library-selective-build.html for more information."
     exit 0