gemini-code-assist[bot] commented on code in PR #19812:
URL: https://github.com/apache/tvm/pull/19812#discussion_r3426774960
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -5033,6 +5071,78 @@ def convert_bitcast(self, op):
return relax.op.memory.view(in_expr, shape=output_shape,
dtype=output_dtype)
+ def convert_rfft2d(self, op):
+ """Convert TFLite RFFT2D for static no-padding/no-truncation shapes."""
+ from tflite.TensorType import TensorType
+
+ input_tensors = self.get_input_tensors(op)
+ output_tensors = self.get_output_tensors(op)
+ if len(input_tensors) != 2 or len(output_tensors) != 1:
+ raise tvm.error.OpNotImplemented("RFFT2D expects two inputs and
one output")
+
+ data_tensor, fft_length_tensor = input_tensors
+ output_tensor = output_tensors[0]
+ if data_tensor.tensor.Type() != TensorType.FLOAT32:
+ raise tvm.error.OpNotImplemented("RFFT2D input must be float32")
+ if not self._is_tflite_complex64_type(output_tensor.tensor.Type()):
+ raise tvm.error.OpNotImplemented("RFFT2D output must be COMPLEX64")
+ if data_tensor.tensor.Sparsity() is not None or
fft_length_tensor.tensor.Sparsity() is not None:
+ raise tvm.error.OpNotImplemented("RFFT2D does not support sparse
inputs")
+
+ input_shape = tuple(to_int_list(self.get_tensor_shape(data_tensor)))
+ tflite_output_shape =
tuple(to_int_list(self.get_tensor_shape(output_tensor)))
+ if len(input_shape) < 2:
+ raise tvm.error.OpNotImplemented("RFFT2D input rank must be at
least 2")
+
+ try:
+ fft_length_value =
self.get_tensor_value_or_prefetched(fft_length_tensor)
+ except (ValueError, TypeError):
+ raise tvm.error.OpNotImplemented("RFFT2D requires a constant
fft_length") from None
Review Comment:
If `fft_length_tensor` is not constant (e.g., a dynamic input),
`get_tensor_value_or_prefetched` can return `None`. Accessing `.dtype` on
`None` would raise an unhandled `AttributeError`. Adding a defensive check for
`None` ensures we raise a clean `OpNotImplemented` exception.
```suggestion
try:
fft_length_value =
self.get_tensor_value_or_prefetched(fft_length_tensor)
except (ValueError, TypeError):
raise tvm.error.OpNotImplemented("RFFT2D requires a constant
fft_length") from None
if fft_length_value is None:
raise tvm.error.OpNotImplemented("RFFT2D requires a constant
fft_length")
```
##########
python/tvm/relax/frontend/tflite/tflite_frontend.py:
##########
@@ -7679,6 +7797,302 @@ def get_tensor_shape(self, tensor_wrapper):
)
+def _is_power_of_2(n):
+ """Return True iff ``n`` is a positive power of 2."""
+ return n > 0 and (n & (n - 1)) == 0
+
+
+def _bit_reversal_swap_pairs(n):
+ """Return the (i, j) index pairs (i < j) for the bit-reversal permutation
of length n.
+
+ For a Cooley-Tukey radix-2 FFT, the input must be permuted by bit-reversing
+ each index in log2(n) bits before the butterfly stages. Precomputing the
+ swap pairs as constants is much cheaper in TIR than computing the
+ bit-reverse on the fly.
+ """
+ assert _is_power_of_2(n), f"bit-reversal requires power of 2, got {n}"
+ length = n.bit_length() - 1 # log2(n)
+ swaps = []
+ for i in range(1, n):
+ j = 0
+ for k in range(length):
+ if i & (1 << k):
+ j |= 1 << (length - 1 - k)
+ if i < j:
+ swaps.append((i, j))
+ return swaps
+
+
+def _build_tflite_rfft2d_primfunc(input_shape, output_pair_shape):
+ """Build a reference TIR kernel for TFLite RFFT2D.
+
+ The TFLite frontend represents complex tensors as float32 real/imag pairs
+ with a trailing dimension of size 2 because TVM does not have a native
+ complex64 dtype. This kernel computes the unnormalized 2-D real FFT over
+ the last two input dimensions and writes that pair representation.
+
+ All trig and accumulation are in float32, so the result agrees with
+ ``np.fft.rfft2`` to about ``1e-5`` absolute tolerance for typical input
+ sizes. Higher-precision backends should override this kernel.
+
+ Notes
+ -----
+ This is a **naive O(B * H * W * H * W) DFT**, not an FFT. For an input of
+ spatial shape (H, W) the inner sum runs H*W times per output position, and
+ there are H*W' output positions per batch (W' = W // 2 + 1). This is
+ intentionally simple for correctness validation against
+ ``np.fft.rfft2``; production use cases with large spatial dimensions should
+ override the kernel with an FFT-based implementation. The outer
+ (batch, out_y, out_x) iteration is structured as S-TIR spatial axes so a
+ downstream ``tvm.tir.schedule`` pass can parallelize it.
+ """
+ from tvm.script.parser import tirx as T
+
+ batch = 1
+ for dim in input_shape[:-2]:
+ batch *= int(dim)
+ height = int(input_shape[-2])
+ width = int(input_shape[-1])
+ out_width = int(output_pair_shape[-2])
+ input_total = batch * height * width
+ output_complex_total = batch * height * out_width
+ neg_two_pi = np.float32(-2.0 * math.pi)
+
+ @T.prim_func(private=True, s_tir=True, check_well_formed=False)
+ def kernel(
+ data: T.Buffer(input_shape, "float32"), output:
T.Buffer(output_pair_shape, "float32")
+ ):
+ # Flat 1D aliases of the multi-dim buffers. The kernel is rank-agnostic
+ # over the leading batch dimensions, so collapsing the index space
+ # avoids special-casing 2D / 3D / 4D input shapes.
+ data_flat = T.decl_buffer((input_total,), "float32", data=data.data)
+ output_flat = T.decl_buffer((output_complex_total * 2,), "float32",
data=output.data)
+ neg_two_pi_const = T.float32(neg_two_pi)
+
+ for b_idx, out_y, out_x in T.grid(batch, height, out_width):
+ with T.sblock("rfft2d"):
+ v_b, v_oy, v_ox = T.axis.remap("SSS", [b_idx, out_y, out_x])
+ real_sum = T.float32(0)
+ imag_sum = T.float32(0)
+ input_base = v_b * height * width
+ for in_y, in_x in T.grid(height, width):
+ phase_y = T.Cast("float32", v_oy) * T.Cast("float32",
in_y) / T.float32(
+ height
+ )
+ phase_x = T.Cast("float32", v_ox) * T.Cast("float32",
in_x) / T.float32(width)
+ angle = neg_two_pi_const * (phase_y + phase_x)
+ value = data_flat[input_base + in_y * width + in_x]
+ real_sum = real_sum + value * T.cos(angle)
+ imag_sum = imag_sum + value * T.sin(angle)
+ flat_out_idx = ((v_b * height + v_oy) * out_width + v_ox) * 2
+ output_flat[flat_out_idx] = real_sum
+ output_flat[flat_out_idx + 1] = imag_sum
+
+ return kernel
+
+
+def _build_tflite_rfft2d_fft_primfunc(input_shape, output_pair_shape):
+ """Build a 2D Cooley-Tukey FFT TIR kernel for TFLite RFFT2D.
+
+ Precondition: both ``input_shape[-2]`` (height) and ``input_shape[-1]``
+ (width) must be positive powers of 2. The frontend dispatches to this
+ kernel via ``_is_power_of_2`` checks; the DFT reference kernel handles
+ the remaining cases (odd / non-power-of-2 sizes).
+
+ Algorithm
+ ---------
+ 1. Copy the real input into a scratch complex buffer (imag = 0) of shape
+ ``(B * H * W,)``.
+ 2. For each batch and each row, run an in-place radix-2 1D FFT of length
+ ``W`` along the width axis.
+ 3. For each batch and each column, run an in-place radix-2 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 permutation required by iterative Cooley-Tukey is done
+ by emitting the (i, j) swap pairs directly in the TIR source (one
+ inlined swap per pair), avoiding the need for runtime index tables.
+
+ Complexity is ``O(B * H * W * (log2(H) + log2(W)))``, vs the DFT
+ reference kernel's ``O(B * H * W * H * W)``.
+ """
+ from tvm.script.parser import tirx as T
+
+ batch = 1
+ for dim in input_shape[:-2]:
+ batch *= int(dim)
+ height = int(input_shape[-2])
+ width = int(input_shape[-1])
+ out_width = int(output_pair_shape[-2])
+ input_total = batch * height * width
+ output_complex_total = batch * height * out_width
+ # Cast to Python float so the f-string-interpolated repr is a plain number
+ # (np.float32's repr is "np.float32(...)", which the TIR parser can't
resolve).
+ neg_two_pi = float(np.float32(-2.0 * math.pi))
+ log2_w = int(math.log2(width))
+ log2_h = int(math.log2(height))
+
+ if not (_is_power_of_2(height) and _is_power_of_2(width)):
+ raise ValueError(
+ f"_build_tflite_rfft2d_fft_primfunc requires power-of-2 height and
width, "
+ f"got H={height}, W={width}"
+ )
+
+ # Precompute the bit-reversal swap pairs at Python level. These are
+ # constant for a given FFT length and will be inlined in the TIR source.
+ # Each emitted line is indented 16 spaces (4 levels: top → b_idx loop →
+ # sblock → row/col loop body) so it lands inside the for loop when
+ # concatenated into the primfunc source.
+ row_swap_stmts = []
+ for i, j in _bit_reversal_swap_pairs(width):
+ row_swap_stmts.append(
+ f" i_idx = row_base + {i}\n"
+ f" j_idx = row_base + {j}\n"
+ f" tmp_r = scratch_real[i_idx]\n"
+ f" scratch_real[i_idx] = scratch_real[j_idx]\n"
+ f" scratch_real[j_idx] = tmp_r\n"
+ f" tmp_i = scratch_imag[i_idx]\n"
+ f" scratch_imag[i_idx] = scratch_imag[j_idx]\n"
+ f" scratch_imag[j_idx] = tmp_i\n"
+ )
+ row_swaps_code = "".join(row_swap_stmts) if row_swap_stmts else "
pass\n"
+
+ col_swap_stmts = []
+ for i, j in _bit_reversal_swap_pairs(height):
+ col_swap_stmts.append(
+ f" i_idx = col_base + {i * width}\n"
+ f" j_idx = col_base + {j * width}\n"
+ f" tmp_r = scratch_real[i_idx]\n"
+ f" scratch_real[i_idx] = scratch_real[j_idx]\n"
+ f" scratch_real[j_idx] = tmp_r\n"
+ f" tmp_i = scratch_imag[i_idx]\n"
+ f" scratch_imag[i_idx] = scratch_imag[j_idx]\n"
+ f" scratch_imag[j_idx] = tmp_i\n"
+ )
+ col_swaps_code = "".join(col_swap_stmts) if col_swap_stmts else "
pass\n"
+
+ # Build the per-stage butterfly code with the stage loop fully unrolled
+ # at primfunc-construction time. After unrolling, all loop bounds
+ # (block_start, k) are compile-time integers, so the TIR parser doesn't
+ # need to reason about runtime loop extents and the scheduler can
+ # optimize the trig calls per stage.
+ def _stage_stmts(stage_count, length, indent, stride=1,
base_expr="row_base"):
+ """Generate fully-unrolled Cooley-Tukey butterfly stage bodies.
+
+ ``base_expr`` is the TIR expression holding the base offset of the
+ FFT being transformed (e.g. ``"row_base"`` for rows or
+ ``"col_base"`` for columns). ``stride`` is the integer distance
+ between adjacent butterfly taps: 1 for the row-FFT (contiguous
+ elements) and ``width`` for the column-FFT (strided access).
+ """
+ out = []
+ for stage in range(1, stage_count + 1):
+ m_val = 1 << stage
+ half_val = m_val >> 1
+ for block_start in range(0, length, m_val):
+ for k in range(half_val):
+ if stride == 1:
+ a_idx = f"{base_expr} + {block_start} + {k}"
+ b_idx_expr = f"{base_expr} + {block_start} + {k} +
{half_val}"
+ else:
+ a_idx = f"{base_expr} + ({block_start} + {k}) *
{stride}"
+ b_idx_expr = (
+ f"{base_expr} + ({block_start} + {k} + {half_val})
* {stride}"
+ )
+ out.extend(
+ [
+ f"{indent}angle = neg_two_pi_const *
T.Cast('float32', {k}) / T.Cast('float32', {m_val})\n",
+ f"{indent}w_real = T.cos(angle)\n",
+ f"{indent}w_imag = T.sin(angle)\n",
+ f"{indent}a_idx = {a_idx}\n",
+ f"{indent}b_idx_local = {b_idx_expr}\n",
+ f"{indent}t_real = scratch_real[b_idx_local] *
w_real - scratch_imag[b_idx_local] * w_imag\n",
+ f"{indent}t_imag = scratch_real[b_idx_local] *
w_imag + scratch_imag[b_idx_local] * w_real\n",
+ f"{indent}u_real = scratch_real[a_idx]\n",
+ f"{indent}u_imag = scratch_imag[a_idx]\n",
+ f"{indent}scratch_real[a_idx] = u_real + t_real\n",
+ f"{indent}scratch_imag[a_idx] = u_imag + t_imag\n",
+ f"{indent}scratch_real[b_idx_local] = u_real -
t_real\n",
+ f"{indent}scratch_imag[b_idx_local] = u_imag -
t_imag\n",
+ ]
+ )
Review Comment:
Since the Cooley-Tukey butterfly stages are fully unrolled at compile time,
the twiddle factor angles are completely static. Instead of generating TIR code
that computes the angles and calls `T.cos` and `T.sin` at runtime, we can
precompute `w_real` and `w_imag` directly in Python and inline them as float
literals. This avoids thousands of expensive runtime transcendental function
calls (e.g., 2048 calls for a 256-point FFT) and reduces the size of the
generated TIR code by ~23%, speeding up both compilation and execution.
```suggestion
angle_val = neg_two_pi * k / m_val
w_real_val = math.cos(angle_val)
w_imag_val = math.sin(angle_val)
out.extend(
[
f"{indent}a_idx = {a_idx}\n",
f"{indent}b_idx_local = {b_idx_expr}\n",
f"{indent}t_real = scratch_real[b_idx_local] *
T.float32({w_real_val!r}) - scratch_imag[b_idx_local] *
T.float32({w_imag_val!r})\n",
f"{indent}t_imag = scratch_real[b_idx_local] *
T.float32({w_imag_val!r}) + scratch_imag[b_idx_local] *
T.float32({w_real_val!r})\n",
f"{indent}u_real = scratch_real[a_idx]\n",
f"{indent}u_imag = scratch_imag[a_idx]\n",
f"{indent}scratch_real[a_idx] = u_real +
t_real\n",
f"{indent}scratch_imag[a_idx] = u_imag +
t_imag\n",
f"{indent}scratch_real[b_idx_local] = u_real -
t_real\n",
f"{indent}scratch_imag[b_idx_local] = u_imag -
t_imag\n",
]
)
```
--
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]
