Arm backend: Add TOSA dialect FFT node visitors

Signed-off-by: Saoirse Stewart <saoirse.stewart@arm.com>
Change-Id: I16991695a3f1d8828e4c17b3f1d4b4b32157298d

Author: SaoirseARM

backends/arm/_passes/aten_to_tosa_tensor_operators.py

@@ -26,6 +26,20 @@ def rewrite_argmax(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec:
     )
 
 
+def rewrite_rfft2(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec | None:
+    fft_size = node.args[1] if len(node.args) > 1 else node.kwargs.get("s")
+    fft_dims = node.args[2] if len(node.args) > 2 else node.kwargs.get("dim", [-2, -1])
+    norm = node.args[3] if len(node.args) > 3 else node.kwargs.get("norm")
+    if fft_size is not None or fft_dims not in ([-2, -1], (-2, -1)) or norm is not None:
+        return None
+
+    return DialectNodeSpec(
+        exir_ops.backend.tosa.RFFT2D.default,
+        (node.args[0],),
+        {},
+    )
+
+
 def rewrite_binary_operator(
     node: Node, pass_: AtenToDialectPass
 ) -> DialectNodeSpec | None:

backends/arm/_passes/exir_to_tosa_pass.py

@@ -12,6 +12,7 @@
 from executorch.backends.arm._passes.aten_to_tosa_tensor_operators import (
     rewrite_argmax,
     rewrite_binary_operator,
+    rewrite_rfft2,
 )
 from executorch.backends.transforms.aten_to_dialect_pass import (
     AtenToDialectPass,
@@ -43,13 +44,24 @@ def decorator(func: SubstitutionFn) -> SubstitutionFn:
     return decorator
 
 
-@ExirToTosaPass.register_dialect_substitution(exir_ops.edge.aten.argmax.default)
+@register_dialect_substitutions(
+    exir_ops.edge.aten.argmax.default,
+)
 def _get_tensor_operators_replacement(
     node: Node, pass_: AtenToDialectPass
-) -> DialectNodeSpec:
+) -> DialectNodeSpec | None:
     return rewrite_argmax(node, pass_)
 
 
+@register_dialect_substitutions(
+    exir_ops.edge.aten.fft_rfft2.default,
+)
+def _get_fft_replacement(
+    node: Node, pass_: AtenToDialectPass
+) -> DialectNodeSpec | None:
+    return rewrite_rfft2(node, pass_)
+
+
 @register_dialect_substitutions(
     exir_ops.edge.aten.add.Tensor,
     exir_ops.edge.aten.bitwise_and.Tensor,

backends/arm/operators/__init__.py

@@ -46,6 +46,7 @@
     op_tosa_depthwise_conv2d,
     op_tosa_eq,
     op_tosa_erf,
+    op_tosa_fft,
     op_tosa_gather,
     op_tosa_ge,
     op_tosa_gt,

backends/arm/operators/op_tosa_fft.py

@@ -0,0 +1,76 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Any, List
+
+import torch.fx
+import tosa_serializer as ts
+
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+    validate_same_dtype,
+    validate_valid_dtype,
+)
+from executorch.backends.arm.tosa.mapping import TosaArg
+
+
+@register_node_visitor
+class FFT2dVisitor(NodeVisitor):
+    target = "tosa.FFT2D.default"
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        validate_num_inputs(self.target, inputs, 2)
+        validate_same_dtype(self.target, inputs, ts)
+        validate_valid_dtype(self.target, inputs, ts.DType.FP32, self.tosa_spec)
+
+        attr = ts.TosaSerializerAttribute()
+        attr.FFT2dAttribute(
+            node.kwargs.get("inverse", False),
+            node.kwargs.get("local_bound", False),
+        )
+        self._serialize_operator(
+            node,
+            tosa_graph,
+            ts.Op.FFT2D,
+            [inputs[0].name, inputs[1].name],
+            output.multiple_output_names,
+            attr,
+        )
+
+
+@register_node_visitor
+class RFFT2dVisitor(NodeVisitor):
+    target = "tosa.RFFT2D.default"
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        validate_num_inputs(self.target, inputs, 1)
+        validate_valid_dtype(self.target, inputs, ts.DType.FP32, self.tosa_spec)
+
+        attr = ts.TosaSerializerAttribute()
+        attr.RFFT2dAttribute(node.kwargs.get("local_bound", False))
+        self._serialize_operator(
+            node,
+            tosa_graph,
+            ts.Op.RFFT2D,
+            [inputs[0].name],
+            output.multiple_output_names,
+            attr,
+        )

backends/arm/test/ops/test_fft.py

@@ -0,0 +1,62 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Tuple
+
+import torch
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import (
+    TosaPipelineFP,
+    VgfPipeline,
+)
+
+aten_op = "torch.ops.aten.fft_rfft2.default"
+exir_op = "executorch_exir_dialects_edge__ops_aten_fft_rfft2_default"
+
+input_t1 = Tuple[torch.Tensor]
+
+
+class RFFT2D(torch.nn.Module):
+    test_parameters = {
+        "rank2": lambda: (torch.randn(8, 16),),
+        "rank3": lambda: (torch.randn(2, 8, 16),),
+        "rank4": lambda: (torch.randn(1, 2, 8, 16),),
+        "ones": lambda: (torch.ones(2, 8, 16),),
+        "zeros": lambda: (torch.zeros(2, 8, 16),),
+    }
+
+    def forward(self, x: torch.Tensor):
+        output = torch.fft.rfft2(x)
+        return output.real, output.imag
+
+
+@common.parametrize("test_data", RFFT2D.test_parameters)
+def test_rfft2d_tosa_FP(test_data: input_t1):
+    pipeline = TosaPipelineFP[input_t1](
+        RFFT2D(),
+        test_data(),
+        aten_op,
+        exir_op,
+        run_on_tosa_ref_model=False,
+        tosa_version="1.1",
+        tosa_extensions=["fft"],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", RFFT2D.test_parameters)
+@common.SkipIfNoModelConverter
+def test_rfft2d_vgf_no_quant(test_data: input_t1):
+    pipeline = VgfPipeline[input_t1](
+        RFFT2D(),
+        test_data(),
+        aten_op,
+        exir_op,
+        run_on_vulkan_runtime=False,
+        quantize=False,
+        tosa_version="TOSA-1.1+FP",
+        tosa_extensions=["fft"],
+    )
+    pipeline.run()

backends/arm/test/targets.bzl

@@ -37,6 +37,7 @@ def define_arm_tests():
         "ops/test_cos.py",
         "ops/test_to_copy.py",
         "ops/test_exp.py",
+        "ops/test_fft.py",
         "ops/test_reciprocal.py",
         "ops/test_mean_dim.py",
         "ops/test_var.py",