Aharrypotter opened a new pull request, #19812:
URL: https://github.com/apache/tvm/pull/19812
## Summary
This PR adds Relax TFLite frontend support for the TFLite builtin `RFFT2D`
operator (issue #19519 item C — FFT / complex operators). It is the
follow-up to upstream PR #19763, which already merged the `REAL` / `IMAG` /
`COMPLEX_ABS` subset of item C; this PR completes the subset with `RFFT2D`
itself.
`RFFT2D` computes a 2D real FFT over the last two input axes and returns a
real/imag pair tensor of shape `[..., H, W // 2 + 1, 2]`. Relax does not have
a native complex64 dtype, so the pair output is represented as a `float32`
tensor with a trailing axis of size 2, matching the convention PR #19763
established for `REAL` / `IMAG` / `COMPLEX_ABS`. `RFFT2D` accepts a real
`float32` input and emits a `float32` pair output of the same dtype; the
frontend does not need to materialize any COMPLEX64 in-memory representation
of its own.
The supported subset is the static no-padding / no-truncation path:
- the input's last two dimensions must match `fft_length`
- `fft_length` must be a length-2 integer constant (int32 or int64)
- the output shape is `input_shape[:-2] + (H, W // 2 + 1, 2)`
- sparse inputs are rejected
## Design
### Dispatch Between Two TIR Kernels
`convert_rfft2d` selects one of two TIR primfuncs at lowering time, based on
whether the spatial axes are powers of two:
```python
if _is_power_of_2(height) and _is_power_of_2(width):
prim_func = _build_tflite_rfft2d_fft_primfunc(input_shape,
relax_output_shape)
else:
prim_func = _build_tflite_rfft2d_primfunc(input_shape,
relax_output_shape)
```
Both kernels share the same `call_tir` contract, so downstream code is
kernel-agnostic.
#### DFT Reference Kernel (`_build_tflite_rfft2d_primfunc`)
A naive O(B · H · W · H · W) DFT over the last two input axes. The outer
(batch, out_y, out_x) iteration is structured as S-TIR spatial axes so a
downstream `tir.schedule` pass can parallelize it. Trig and accumulation are
in float32; the result agrees with `np.fft.rfft2` to about `1e-5` absolute
tolerance for typical input sizes. This kernel is the fallback for odd or
non-power-of-2 spatial sizes.
#### Cooley-Tukey FFT Kernel (`_build_tflite_rfft2d_fft_primfunc`)
An O(B · H · W · (log2(H) + log2(W))) radix-2 Cooley-Tukey FFT, dispatched
when both `H` and `W` are positive powers of 2. The primfunc source is
generated as a TIR string at construction time and registered in `linecache`
so `tirx.parser` can resolve it. The bit-reversal permutation and butterfly
stages are precomputed in Python and inlined as direct scratch-buffer
assignments, so all loop bounds and twiddle factors are compile-time
literals:
1. Copy the real input into `scratch_real`; initialize `scratch_imag` to 0.
2. For each batch and each row, run an in-place 1D FFT of length `W` along
the width axis using scratch buffers.
3. For each batch and each column, run an in-place 1D FFT of length `H`
along the height axis with stride `W`.
4. Write the first `W // 2 + 1` complex bins per row to the output pair
representation.
The bit-reversal swap pairs and butterfly stage bodies are precomputed at
primfunc-construction time because `tirx.parser` does not currently accept
runtime `T.serial` bounds, and twiddle factors must be `T.float32` literals
to avoid `Undefined variable` errors. The fake linecache filename
`<tflite_rfft2d_fft_primfunc H=8 W=8 outW=5>` is dimension-aware, so
generated-source stack traces are readable.
### COMPLEX64 Pair Representation Helpers
The frontend represents TFLite COMPLEX64 tensors as float32 real/imag pairs
with a trailing axis of size 2 (since Relax has no native complex64 dtype).
Four small helpers centralize the rule so future complex ops can plug in
without re-implementing the pair-axis layout:
- `_is_tflite_complex64_type` — checks whether a TFLite tensor type is
`COMPLEX64`.
- `_unwrap_tflite_tensor` — unwraps a `TensorWrapper` to the raw
`tflite.Tensor`.
- `_get_relax_tensor_dtype` — returns the Relax dtype used to represent a
TFLite tensor (`"float32"` for COMPLEX64, otherwise the standard mapping).
- `_get_relax_tensor_shape` — returns the Relax shape (TFLite shape with a
trailing `(2,)` axis appended for COMPLEX64).
The 3 callers that construct Relax parameters from TFLite metadata
(`_get_static_tensor_shape_dtype` / `_set_subgraph_input_params` /
`_get_tensor_param`) now go through these helpers. The pair-axis invariant
is documented on `get_tensor_value` and `get_tensor_shape`, which both
return the *raw TFLite* shape (no pair axis) for downstream callers that
need to compare against the model's declared output shape.
### Boundary Validation
`convert_rfft2d` validates rank, dtype, fft_length shape, integer-ness,
positivity, fft_length == input spatial shape, output shape agreement, and
the absence of sparse inputs before emitting the `call_tir`. Edge cases
(sparse, non-integer fft_length, zero/negative fft_length, mismatched
fft_length, 1×1 spatial size) each produce a targeted `OpNotImplemented`
diagnostic.
## Operator Support
| Operator | TFLite options | Relax lowering | Supported subset |
|---|---|---|---|
| `RFFT2D` | input `float32`, constant length-2 integer `fft_length`,
`COMPLEX64` output | `call_tir` to a generated TIR kernel | static
no-padding/no-truncation; `H`, `W` arbitrary; Cooley-Tukey dispatched when both
are powers of 2 |
## Safety Checks
- Non-float32 input raises `OpNotImplemented("RFFT2D input must be
float32")`.
- Non-COMPLEX64 output raises `OpNotImplemented("RFFT2D output must be
COMPLEX64")`.
- Sparse inputs raise `OpNotImplemented("RFFT2D does not support sparse
inputs")`.
- Non-constant `fft_length` raises `OpNotImplemented("RFFT2D requires a
constant fft_length")`.
- Non-integer `fft_length` raises `OpNotImplemented("RFFT2D fft_length must
be an integer tensor")`.
- Wrong-shape `fft_length` (not length 2) raises `OpNotImplemented("RFFT2D
fft_length must be a length-2 tensor")`.
- Non-positive `fft_length` raises `OpNotImplemented("RFFT2D fft_length must
be positive")`.
- `fft_length` not matching the input's last two dims raises
`OpNotImplemented("RFFT2D currently supports fft_length matching the input
spatial shape")`.
- Mismatched output shape raises `OpNotImplemented("RFFT2D output shape does
not match fft_length")`.
- Input rank < 2 raises `OpNotImplemented("RFFT2D input rank must be at
least 2")`.
## Not Included
- `RFFT2D` with `fft_length` not matching the input's last two dimensions
(padding / truncation path).
- Other complex-data operators: `REAL` / `IMAG` / `COMPLEX_ABS` are already
handled by upstream PR #19763 and are out of scope for this PR.
- A frontend guard that rejects COMPLEX64 inputs flowing into non-complex
ops. With PR #19763 providing the only other complex ops, models that
contain only `RFFT2D` (with float32 input / COMPLEX64 output) and the
three upstream ops are fully supported. A generic COMPLEX64 guard was
intentionally not added here to keep this PR scoped to `RFFT2D`.
- User-override validation for `shape_dict` / `dtype_dict` on COMPLEX64
inputs. After PR #19763, the frontend no longer auto-appends a
trailing pair axis to user overrides; a user passing the natural TFLite
shape without the pair axis will now fall through to the standard
metadata mismatch path.
- Higher-precision FFT kernel (e.g. SIMD). The float32 twiddle / float32
accumulation paths match `np.fft.rfft2` to `~1e-4` on the Cooley-Tukey
path; large spatial dimensions may need a future optimized lowering or
backend-specific implementation.
## Tests
The tests manually build minimal TFLite flatbuffers, run the frontend, and
compare against `np.fft.rfft2`. Edge cases raise `OpNotImplemented`. The
DFT-path tests use `atol=1e-5`; the FFT-path tests use `atol=1e-4` because
twiddle factors are float32 literals.
| Test | Path | Shape | Coverage |
|---|---|---|---|
| `test_rfft2d_static_pair_output` | DFT | `[2, 4]` | Baseline 2D, even
width; also asserts Relax script contains the `tflite_rfft2d` kernel name and
pair-output struct-info |
| `test_rfft2d_static_pair_output_with_batch` | DFT | `[2, 2, 4]` | Leading
batch dims preserved |
| `test_rfft2d_odd_width_pair_output` | DFT | `[3, 5]` | Odd width → `W // 2
+ 1` output bins |
| `test_rfft2d_int64_fft_length` | DFT | `[2, 4]` | INT64 fft_length
constant (TFLite schema allows either int32 or int64) |
| `test_rfft2d_4d_input_pair_output` | DFT | `[2, 3, 4, 5]` | 4D input with
batch and odd width |
| `test_rfft2d_minimal_1x1_pair_output` | DFT | `[1, 1]` | Edge case:
trivial 1×1 FFT |
| `test_rfft2d_mismatched_fft_length_unsupported` | — | `[2, 4]` (fft=`[4,
4]`) | fft_length != input spatial shape guard |
| `test_rfft2d_fft_path_4x4` | **FFT** | `[4, 4]` | Smallest power-of-two
(4×4) where both row and column FFTs do real work |
| `test_rfft2d_fft_path_8x8` | **FFT** | `[8, 8]` | Square 8×8 power-of-two |
| `test_rfft2d_fft_path_16x16` | **FFT** | `[16, 16]` | Larger FFT, kernel
scaling check |
| `test_rfft2d_fft_path_2x2x4x8` | **FFT** | `[2, 2, 4, 8]` | 4D
power-of-two with batch |
Local validation:
```bash
python -m py_compile \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m ruff check \
python/tvm/relax/frontend/tflite/tflite_frontend.py \
tests/python/relax/test_frontend_tflite.py
python -m pytest \
tests/python/relax/test_frontend_tflite.py \
-k rfft2d -v
```
Result:
```text
py_compile: passed
ruff check: All checks passed
rfft2d tests: 11 passed
```
## References
- Issue [#19519](https://github.com/apache/tvm/issues/19519) item C: FFT /
complex operators (`RFFT2D`, `REAL`, `IMAG`, `COMPLEX_ABS`).
- Upstream PR #19763: `[Relax][Frontend][TFLite] Add support for
FFT/complex operators: REAL, IMAG, COMPLEX_ABS` (merge commit `9d6e1cf0`).
This PR's RFFT2D output pair layout matches the COMPLEX64 pair
representation PR #19763 established for `REAL` / `IMAG` / `COMPLEX_ABS`,
so downstream ops from PR #19763 can consume RFFT2D output directly.
- Tracking issue [#19764](https://github.com/apache/tvm/issues/19764) for
the longer-term native `relax.op.signal.rfft2d` / registered TOPI backend
path. This PR is a frontend-local lowering for TFLite `RFFT2D`, not the
native Relax signal op tracked there.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
