Aharrypotter commented on code in PR #19812:
URL: https://github.com/apache/tvm/pull/19812#discussion_r3426980981
##########
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:
make sense, fixed
--
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]
