Arm backend: Add FP8 support for conv, pool, and matmul (pytorch#20108)

Add TOSA FP8E4M3 and FP8E5M2 lowering support for:

CONV2D, DEPTHWISE_CONV2D, CONV3D, TRANSPOSE_CONV2D,
AVG_POOL2D, MAX_POOL2D, and MATMUL.
    
Use wider TOSA outputs for FP8 convolution and matmul, then cast
back when the exported graph expects an FP8 output.
    
Change-Id: I914b7861dd41061130d7a50797ea58e0fe09a4cd

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

---------

Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>

Author: YufengShi-dudu

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/test/ops/test_adaptive_avg_pool2d.py

@@ -112,6 +112,19 @@ def forward(self, *args, **kwargs):
     ),
 }
 
+test_modules_fp8 = {
+    "output_2x2_fp8e4m3": lambda: (
+        AdaptiveAvgPool2d((2, 2)),
+        (torch.rand(1, 4, 10, 10).to(torch.float8_e4m3fn),),
+        "fp8e4m3",
+    ),
+    "output_2x2_fp8e5m2": lambda: (
+        AdaptiveAvgPool2d((2, 2)),
+        (torch.rand(1, 4, 10, 10).to(torch.float8_e5m2),),
+        "fp8e5m2",
+    ),
+}
+
 
 @common.parametrize("test_module", test_modules)
 def test_adaptive_avg_pool2d_tosa_FP(test_module):
@@ -126,6 +139,22 @@ def test_adaptive_avg_pool2d_tosa_FP(test_module):
     pipeline.run()
 
 
+@common.parametrize("test_module", test_modules_fp8)
+def test_adaptive_avg_pool2d_tosa_FP_fp8(test_module):
+    model, input_tensor, tosa_extension = test_module()
+
+    pipeline = TosaPipelineFP[input_t](
+        model,
+        input_tensor,
+        aten_op=[],
+        exir_op=exir_op,
+        tosa_extensions=[tosa_extension],
+        run_on_tosa_ref_model=False,  # torch.avg_pool2d() has no eager CPU FP8 implementation, so eager reference execution fails.
+    )
+    pipeline.count_tosa_ops({"AVG_POOL2D": 4})
+    pipeline.run()
+
+
 @common.parametrize("test_module", test_modules)
 def test_adaptive_avg_pool2d_tosa_INT(test_module):
     model, input_tensor = test_module()

backends/arm/test/ops/test_avg_pool2d.py

@@ -150,6 +150,18 @@ def forward(self, x: torch.Tensor):
         (torch.rand(1, 4, 12, 12, dtype=torch.float16),),
     ),
 }
+test_modules_fp8 = {
+    "rand_fp8e4m3": lambda: (
+        AvgPool2d(4, 2, 0, False),
+        (torch.rand(1, 16, 50, 32).to(torch.float8_e4m3fn),),
+        "fp8e4m3",
+    ),
+    "kernel_3x3_stride_1_pad_1_fp8e5m2": lambda: (
+        AvgPool2d((3, 3), (1, 1), 1),
+        (torch.rand(1, 4, 12, 12).to(torch.float8_e5m2),),
+        "fp8e5m2",
+    ),
+}
 
 
 @common.parametrize("test_module", test_modules | test_modules_bf16 | test_modules_fp16)
@@ -166,6 +178,21 @@ def test_avg_pool2d_tosa_FP(test_module):
     pipeline.run()
 
 
+@common.parametrize("test_module", test_modules_fp8)
+def test_avg_pool2d_tosa_FP_fp8(test_module):
+    model, input_tensor, tosa_extension = test_module()
+    pipeline = TosaPipelineFP[input_t](
+        model,
+        input_tensor,
+        aten_op,
+        exir_op,
+        tosa_extensions=[tosa_extension],
+        run_on_tosa_ref_model=False,  # torch.avg_pool2d() has no eager CPU FP8 implementation, so eager reference execution fails.
+    )
+    pipeline.count_tosa_ops({"AVG_POOL2D": 1})
+    pipeline.run()
+
+
 @common.parametrize("test_module", test_modules)
 def test_avg_pool2d_tosa_INT(test_module):
     model, input_tensor = test_module()