Skip to content

Commit fe9bc95

Browse files
SaoirseARMSaoirseARM
authored
Arm backend: Add TOSA dialect FFT ops (pytorch#20111)
Added TOSA dialect for: - FFT2D - RFFT2D Signed-off-by: Saoirse Stewart <saoirse.stewart@arm.com>
1 parent b36bb84 commit fe9bc95

4 files changed

Lines changed: 276 additions & 0 deletions

File tree

‎backends/arm/test/misc/test_tosa_dialect_fft.py‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
import executorch.backends.arm.tosa.dialect # noqa: F401
7+
import pytest
8+
import sympy # type: ignore[import-untyped]
9+
import torch
10+
from executorch.backends.arm.tosa.dialect.lib import TosaValueError
11+
from executorch.backends.arm.tosa.specification import (
12+
TosaLoweringContext,
13+
TosaSpecification,
14+
)
15+
from executorch.exir.dialects._ops import ops as exir_ops
16+
from torch._subclasses.fake_tensor import FakeTensorMode
17+
from torch.fx.experimental.symbolic_shapes import ShapeEnv
18+
19+
20+
def _make_symint(
21+
shape_env: ShapeEnv, symbol: str, hint: int, min: int = 1, max: int = 64
22+
) -> torch.SymInt:
23+
symint = shape_env.create_symintnode(sympy.Symbol(symbol), hint=hint)
24+
assert isinstance(symint, torch.SymInt)
25+
shape_env.constrain_symbol_range(
26+
symint.node.expr, compiler_min=min, compiler_max=max
27+
)
28+
return symint
29+
30+
31+
def _expr(sym: torch.SymInt) -> sympy.Expr:
32+
return sympy.sympify(str(sym.node._expr))
33+
34+
35+
def test_fft2d_tosa_fp_fft() -> None:
36+
input_real = torch.randn((2, 8, 16), dtype=torch.float32)
37+
input_imag = torch.randn((2, 8, 16), dtype=torch.float32)
38+
39+
with TosaLoweringContext(
40+
TosaSpecification.create_from_string("TOSA-1.1+FP+fft")
41+
), FakeTensorMode() as mode:
42+
output_real, output_imag = exir_ops.backend.tosa.FFT2D.default(
43+
mode.from_tensor(input_real),
44+
mode.from_tensor(input_imag),
45+
)
46+
47+
assert output_real.dtype == torch.float32
48+
assert output_imag.dtype == torch.float32
49+
assert tuple(output_real.shape) == (2, 8, 16)
50+
assert tuple(output_imag.shape) == (2, 8, 16)
51+
52+
53+
def test_fft2d_accepts_matching_symbolic_shape() -> None:
54+
shape_env = ShapeEnv()
55+
width = _make_symint(shape_env, "w", hint=16)
56+
57+
with TosaLoweringContext(
58+
TosaSpecification.create_from_string("TOSA-1.1+FP+fft"),
59+
shape_env,
60+
), FakeTensorMode(shape_env=shape_env) as mode:
61+
input_real = torch.empty((2, 8, width), dtype=torch.float32)
62+
input_imag = torch.empty((2, 8, width), dtype=torch.float32)
63+
output_real, output_imag = exir_ops.backend.tosa.FFT2D.default(
64+
mode.from_tensor(input_real),
65+
mode.from_tensor(input_imag),
66+
)
67+
68+
assert isinstance(output_real.shape[2], torch.SymInt)
69+
assert isinstance(output_imag.shape[2], torch.SymInt)
70+
assert sympy.simplify(_expr(output_real.shape[2]) - sympy.Symbol("w")) == 0
71+
assert sympy.simplify(_expr(output_imag.shape[2]) - sympy.Symbol("w")) == 0
72+
73+
74+
def test_rfft2d_tosa_fp_fft() -> None:
75+
input_real = torch.randn((2, 8, 16), dtype=torch.float32)
76+
77+
with TosaLoweringContext(
78+
TosaSpecification.create_from_string("TOSA-1.1+FP+fft")
79+
), FakeTensorMode() as mode:
80+
output_real, output_imag = exir_ops.backend.tosa.RFFT2D.default(
81+
mode.from_tensor(input_real),
82+
)
83+
84+
assert output_real.dtype == torch.float32
85+
assert output_imag.dtype == torch.float32
86+
assert tuple(output_real.shape) == (2, 8, 9)
87+
assert tuple(output_imag.shape) == (2, 8, 9)
88+
89+
90+
def test_fft_requires_extension() -> None:
91+
input_real = torch.randn((2, 8, 16), dtype=torch.float32)
92+
input_imag = torch.randn((2, 8, 16), dtype=torch.float32)
93+
94+
with TosaLoweringContext(
95+
TosaSpecification.create_from_string("TOSA-1.1+FP")
96+
), FakeTensorMode() as mode:
97+
with pytest.raises(TosaValueError, match="doesn't support FFT2D"):
98+
exir_ops.backend.tosa.FFT2D.default(
99+
mode.from_tensor(input_real),
100+
mode.from_tensor(input_imag),
101+
)
102+
103+
104+
def test_rfft2d_preserves_symbolic_width() -> None:
105+
shape_env = ShapeEnv()
106+
width = _make_symint(shape_env, "w", hint=16)
107+
108+
with TosaLoweringContext(
109+
TosaSpecification.create_from_string("TOSA-1.1+FP+fft"),
110+
shape_env,
111+
), FakeTensorMode(shape_env=shape_env) as mode:
112+
input_real = torch.empty((2, 8, width), dtype=torch.float32)
113+
output_real, output_imag = exir_ops.backend.tosa.RFFT2D.default(
114+
mode.from_tensor(input_real)
115+
)
116+
117+
expected = sympy.floor(sympy.Symbol("w") / 2) + sympy.Integer(1)
118+
assert isinstance(output_real.shape[2], torch.SymInt)
119+
assert isinstance(output_imag.shape[2], torch.SymInt)
120+
assert sympy.simplify(_expr(output_real.shape[2]) - expected) == 0
121+
assert sympy.simplify(_expr(output_imag.shape[2]) - expected) == 0

‎backends/arm/tosa/dialect/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
conv3d,
1212
custom,
1313
depthwise_conv2d,
14+
fft,
1415
gather,
1516
identity,
1617
matmul,

‎backends/arm/tosa/dialect/ops/_common.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
# This source code is licensed under the BSD-style license found in the
44
# LICENSE file in the root directory of this source tree.
55

6+
import torch
67
from executorch.backends.arm.tosa.dialect.lib import TosaValueError
78

89
_VALID_NAN_MODES = {"PROPAGATE", "IGNORE"}
@@ -14,3 +15,12 @@ def validate_nan_mode(nan_mode: str, op: str) -> None:
1415
f"Unsupported nan_mode {nan_mode}. Expected one of {_VALID_NAN_MODES}",
1516
op=op,
1617
)
18+
19+
20+
def validate_power_of_two(size: int | torch.SymInt, name: str, op: str) -> None:
21+
if not isinstance(size, int):
22+
return
23+
if size < 1 or (size & (size - 1)) != 0:
24+
raise TosaValueError(
25+
f"{name} must be a positive power of two, got {size}", op=op
26+
)

‎backends/arm/tosa/dialect/ops/fft.py‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
import sympy # type: ignore[import-untyped]
7+
import torch
8+
from executorch.backends.arm.tosa.dialect.lib import TosaValueError
9+
from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op
10+
from executorch.backends.arm.tosa.specification import (
11+
get_context_shape_env,
12+
get_context_spec,
13+
TosaSpecification,
14+
)
15+
from torch.utils._sympy.functions import FloorDiv
16+
17+
18+
def _validate_fft_spec(op: str) -> None:
19+
tosa_spec = get_context_spec()
20+
if not (tosa_spec.support_float() and tosa_spec.support_extension("fft")):
21+
raise TosaValueError(
22+
f"TOSA spec {tosa_spec} doesn't support {op}",
23+
op=op,
24+
)
25+
26+
27+
def _is_power_of_two(value: int) -> bool:
28+
return value > 0 and (value & (value - 1)) == 0
29+
30+
31+
def _validate_power_of_two(value: int | torch.SymInt, name: str, op: str) -> None:
32+
if isinstance(value, torch.SymInt):
33+
expr = sympy.simplify(_to_sympy_expr(value))
34+
value_range = get_context_shape_env().bound_sympy(expr)
35+
if value_range.is_int and value_range.is_singleton():
36+
singleton = sympy.simplify(value_range.lower)
37+
if singleton.is_integer and not _is_power_of_two(int(singleton)):
38+
raise TosaValueError(
39+
f"{op} requires {name} to be a power of two but got {singleton}",
40+
op=op,
41+
)
42+
return
43+
44+
if not _is_power_of_two(int(value)):
45+
raise TosaValueError(
46+
f"{op} requires {name} to be a power of two but got {value}",
47+
op=op,
48+
)
49+
50+
51+
def _validate_fft_input(input_real: torch.Tensor, op: str) -> None:
52+
if input_real.dtype != torch.float32:
53+
raise TosaValueError(f"{op} requires float32 inputs", op=op)
54+
if input_real.dim() != 3:
55+
raise TosaValueError(f"{op} requires a rank-3 input", op=op)
56+
57+
_, height, width = input_real.shape
58+
_validate_power_of_two(height, "height", op)
59+
_validate_power_of_two(width, "width", op)
60+
61+
62+
def _to_sympy_expr(value: int | torch.SymInt) -> sympy.Expr:
63+
if isinstance(value, torch.SymInt):
64+
return value.node._expr
65+
return sympy.Integer(int(value))
66+
67+
68+
def _rfft_output_width(width: int | torch.SymInt) -> int | torch.SymInt:
69+
if isinstance(width, torch.SymInt):
70+
expr = FloorDiv(_to_sympy_expr(width), sympy.Integer(2)) + sympy.Integer(1)
71+
return get_context_shape_env().create_symintnode(expr, hint=None)
72+
return width // 2 + 1
73+
74+
75+
def _same_fft_dimension(lhs: int | torch.SymInt, rhs: int | torch.SymInt) -> bool:
76+
if not isinstance(lhs, torch.SymInt) and not isinstance(rhs, torch.SymInt):
77+
return lhs == rhs
78+
79+
diff = sympy.simplify(_to_sympy_expr(lhs) - _to_sympy_expr(rhs))
80+
if diff == 0:
81+
return True
82+
83+
value_range = get_context_shape_env().bound_sympy(diff)
84+
return (
85+
value_range.is_int
86+
and value_range.is_singleton()
87+
and sympy.simplify(value_range.lower) == 0
88+
)
89+
90+
91+
def _same_fft_shape(
92+
lhs: torch.Size | tuple[int | torch.SymInt, ...],
93+
rhs: torch.Size | tuple[int | torch.SymInt, ...],
94+
) -> bool:
95+
return len(lhs) == len(rhs) and all(
96+
_same_fft_dimension(lhs_dim, rhs_dim) for lhs_dim, rhs_dim in zip(lhs, rhs)
97+
)
98+
99+
100+
@register_fake_tosa_op(
101+
"FFT2D(Tensor input_real, Tensor input_imag, *, bool inverse=False, bool local_bound=False) -> (Tensor output_real, Tensor output_imag)",
102+
TosaSpecification.all_versions_and_profiles(),
103+
)
104+
def FFT2D(
105+
input_real: torch.Tensor,
106+
input_imag: torch.Tensor,
107+
*,
108+
inverse: bool = False,
109+
local_bound: bool = False,
110+
) -> tuple[torch.Tensor, torch.Tensor]:
111+
_validate_fft_spec("FFT2D")
112+
_validate_fft_input(input_real, "FFT2D")
113+
_validate_fft_input(input_imag, "FFT2D")
114+
115+
if not _same_fft_shape(input_real.shape, input_imag.shape):
116+
raise TosaValueError(
117+
f"FFT2D expects matching input shapes but got {tuple(input_real.shape)} and {tuple(input_imag.shape)}",
118+
op="FFT2D",
119+
)
120+
121+
return (
122+
torch.empty_like(input_real, dtype=input_real.dtype),
123+
torch.empty_like(input_imag, dtype=input_imag.dtype),
124+
)
125+
126+
127+
@register_fake_tosa_op(
128+
"RFFT2D(Tensor input_real, *, bool local_bound=False) -> (Tensor output_real, Tensor output_imag)",
129+
TosaSpecification.all_versions_and_profiles(),
130+
)
131+
def RFFT2D(
132+
input_real: torch.Tensor,
133+
*,
134+
local_bound: bool = False,
135+
) -> tuple[torch.Tensor, torch.Tensor]:
136+
_validate_fft_spec("RFFT2D")
137+
_validate_fft_input(input_real, "RFFT2D")
138+
139+
batch, height, width = input_real.shape
140+
output_shape = (batch, height, _rfft_output_width(width))
141+
return (
142+
torch.empty(output_shape, dtype=input_real.dtype),
143+
torch.empty(output_shape, dtype=input_real.dtype),
144+
)

0 commit comments

Comments
 (0)