build/torch-neuron/__init__.py

"""NKI-accelerated 2D Fourier Transform kernel for FNet on AWS Neuron (Trainium/Inferentia).

Replaces the `FNetBasicFourierTransform` module with NKI kernels that perform
the 2D FFT (over sequence and hidden dimensions) entirely on NeuronCore hardware.

Algorithm:
    FNet applies fftn(x, dim=(1,2)).real to the hidden states tensor (B, S, D).
    This is decomposed into:
    1. 1D FFT along hidden_dim (dim=2) for each (batch, seq) position
    2. 1D FFT along seq_dim (dim=1) for each (batch, hidden) position
    3. Take real part

    The NKI kernel uses radix-2 Cooley-Tukey with a 128-point DFT base case
    computed via nc_matmul on the Tensor Engine (~90% utilization in base case).

Supported configurations:
    - FNet-base: hidden_size=768 (padded to 1024), max_seq_len=512
    - FNet-large: hidden_size=1024, max_seq_len=512
    - Custom configs where hidden_size and seq_len are powers of 2 (128-2048)

Requirements:
    - PyTorch Native (torch-neuronx) with NKI support
    - NKI >= 0.3.0 (SDK 2.29+)
    - hidden_size must be <= 2048 (padded to next power of 2)
    - seq_len must be <= 2048 (padded to next power of 2)
"""

import math

import numpy as np
import torch
import torch.nn as nn

import nki
import nki.isa as nisa
import nki.language as nl


# =============================================================================
# NKI Kernel: 1D FFT (flat radix-2 Cooley-Tukey)
# =============================================================================

TILE_H = 128  # NeuronCore partition height
BASE_N = 128  # Base DFT size (matches Tensor Engine 128x128 systolic array)


def _fft1d_matmul_isa(X_real, X_imag, Y_real, Y_imag, W_real, W_imag):
    """128-pt DFT via complex matrix multiply on the Tensor Engine."""
    H, W_size = X_real.shape

    W_real_T_psum = nl.ndarray((W_size, W_size), dtype=nl.float32, buffer=nl.psum)
    nisa.nc_transpose(dst=W_real_T_psum, data=W_real)
    W_real_T = nl.ndarray((W_size, W_size), dtype=nl.float32, buffer=nl.sbuf)
    nisa.tensor_copy(dst=W_real_T, src=W_real_T_psum)

    W_imag_T_psum = nl.ndarray((W_size, W_size), dtype=nl.float32, buffer=nl.psum)
    nisa.nc_transpose(dst=W_imag_T_psum, data=W_imag)
    W_imag_T = nl.ndarray((W_size, W_size), dtype=nl.float32, buffer=nl.sbuf)
    nisa.tensor_copy(dst=W_imag_T, src=W_imag_T_psum)

    neg_W_imag_T = nl.ndarray((W_size, W_size), dtype=nl.float32, buffer=nl.sbuf)
    nisa.tensor_scalar(dst=neg_W_imag_T, data=W_imag_T, op0=nl.multiply, operand0=-1.0)

    X_real_T_psum = nl.ndarray((W_size, H), dtype=nl.float32, buffer=nl.psum)
    nisa.nc_transpose(dst=X_real_T_psum, data=X_real)
    X_real_T = nl.ndarray((W_size, H), dtype=nl.float32, buffer=nl.sbuf)
    nisa.tensor_copy(dst=X_real_T, src=X_real_T_psum)

    X_imag_T_psum = nl.ndarray((W_size, H), dtype=nl.float32, buffer=nl.psum)
    nisa.nc_transpose(dst=X_imag_T_psum, data=X_imag)
    X_imag_T = nl.ndarray((W_size, H), dtype=nl.float32, buffer=nl.sbuf)
    nisa.tensor_copy(dst=X_imag_T, src=X_imag_T_psum)

    psum_real = nl.ndarray((H, W_size), dtype=nl.float32, buffer=nl.psum)
    nisa.nc_matmul(dst=psum_real, stationary=X_real_T, moving=W_real_T)
    nisa.nc_matmul(dst=psum_real, stationary=X_imag_T, moving=neg_W_imag_T)
    nisa.tensor_copy(dst=Y_real, src=psum_real)

    psum_imag = nl.ndarray((H, W_size), dtype=nl.float32, buffer=nl.psum)
    nisa.nc_matmul(dst=psum_imag, stationary=X_real_T, moving=W_imag_T)
    nisa.nc_matmul(dst=psum_imag, stationary=X_imag_T, moving=W_real_T)
    nisa.tensor_copy(dst=Y_imag, src=psum_imag)


def _complex_twiddle_multiply(
    odd_r, odd_i, tw_r, tw_i, out_r, out_i, tmp_ac, tmp_bd, tmp_ad, tmp_bc
):
    """Complex multiplication: out = twiddle * odd."""
    nisa.tensor_tensor(dst=tmp_ac, data1=odd_r, data2=tw_r, op=nl.multiply)
    nisa.tensor_tensor(dst=tmp_bd, data1=odd_i, data2=tw_i, op=nl.multiply)
    nisa.tensor_tensor(dst=tmp_ad, data1=odd_r, data2=tw_i, op=nl.multiply)
    nisa.tensor_tensor(dst=tmp_bc, data1=odd_i, data2=tw_r, op=nl.multiply)
    nisa.tensor_tensor(dst=out_r, data1=tmp_ac, data2=tmp_bd, op=nl.subtract)
    nisa.tensor_tensor(dst=out_i, data1=tmp_ad, data2=tmp_bc, op=nl.add)


def _butterfly_pair(
    even_r,
    even_i,
    odd_r,
    odd_i,
    tw_r,
    tw_i,
    new_even_r,
    new_even_i,
    new_odd_r,
    new_odd_i,
    tw_out_r,
    tw_out_i,
    tmp_ac,
    tmp_bd,
    tmp_ad,
    tmp_bc,
):
    """Butterfly: new_even = even + twiddle*odd, new_odd = even - twiddle*odd."""
    _complex_twiddle_multiply(
        odd_r, odd_i, tw_r, tw_i, tw_out_r, tw_out_i, tmp_ac, tmp_bd, tmp_ad, tmp_bc
    )
    nisa.tensor_tensor(dst=new_even_r, data1=even_r, data2=tw_out_r, op=nl.add)
    nisa.tensor_tensor(dst=new_even_i, data1=even_i, data2=tw_out_i, op=nl.add)
    nisa.tensor_tensor(dst=new_odd_r, data1=even_r, data2=tw_out_r, op=nl.subtract)
    nisa.tensor_tensor(dst=new_odd_i, data1=even_i, data2=tw_out_i, op=nl.subtract)


# ---------------------------------------------------------------------------
# 128-pt NKI FFT kernel
# ---------------------------------------------------------------------------


@nki.jit
def _fft1d_128(X_real_hbm, X_imag_hbm, W_128_real_hbm, W_128_imag_hbm):
    """128-point FFT: direct DFT via Tensor Engine matmul."""
    Y_real_hbm = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.shared_hbm)
    Y_imag_hbm = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.shared_hbm)

    X_real = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    X_imag = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    nisa.dma_copy(dst=X_real, src=X_real_hbm[0:TILE_H, 0:BASE_N])
    nisa.dma_copy(dst=X_imag, src=X_imag_hbm[0:TILE_H, 0:BASE_N])

    W_real = nl.ndarray((BASE_N, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    W_imag = nl.ndarray((BASE_N, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    nisa.dma_copy(dst=W_real, src=W_128_real_hbm[0:BASE_N, 0:BASE_N])
    nisa.dma_copy(dst=W_imag, src=W_128_imag_hbm[0:BASE_N, 0:BASE_N])

    Y_real = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    Y_imag = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    _fft1d_matmul_isa(X_real, X_imag, Y_real, Y_imag, W_real, W_imag)

    nisa.dma_copy(dst=Y_real_hbm[0:TILE_H, 0:BASE_N], src=Y_real)
    nisa.dma_copy(dst=Y_imag_hbm[0:TILE_H, 0:BASE_N], src=Y_imag)
    return Y_real_hbm, Y_imag_hbm


# ---------------------------------------------------------------------------
# 256-pt NKI FFT kernel
# ---------------------------------------------------------------------------


@nki.jit
def _fft1d_256(
    X_real_hbm,
    X_imag_hbm,
    W_128_real_hbm,
    W_128_imag_hbm,
    twiddle_real_hbm,
    twiddle_imag_hbm,
):
    """256-point FFT: 2 groups of 128, 1 butterfly level."""
    N = 256
    N_half = 128

    Y_real_hbm = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.shared_hbm)
    Y_imag_hbm = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.shared_hbm)

    X_real = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.sbuf)
    X_imag = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.sbuf)
    nisa.dma_copy(dst=X_real, src=X_real_hbm[0:TILE_H, 0:N])
    nisa.dma_copy(dst=X_imag, src=X_imag_hbm[0:TILE_H, 0:N])

    W_128_real = nl.ndarray((BASE_N, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    W_128_imag = nl.ndarray((BASE_N, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    nisa.dma_copy(dst=W_128_real, src=W_128_real_hbm[0:BASE_N, 0:BASE_N])
    nisa.dma_copy(dst=W_128_imag, src=W_128_imag_hbm[0:BASE_N, 0:BASE_N])

    twiddle_real = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    twiddle_imag = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    nisa.dma_copy(dst=twiddle_real, src=twiddle_real_hbm[0:TILE_H, 0:N_half])
    nisa.dma_copy(dst=twiddle_imag, src=twiddle_imag_hbm[0:TILE_H, 0:N_half])

    # Deinterleave
    even_idx = nl.ndarray((TILE_H, N_half), dtype=nl.uint32, buffer=nl.sbuf)
    odd_idx = nl.ndarray((TILE_H, N_half), dtype=nl.uint32, buffer=nl.sbuf)
    nisa.iota(dst=even_idx, pattern=[[2, N_half]], offset=0, channel_multiplier=0)
    nisa.iota(dst=odd_idx, pattern=[[2, N_half]], offset=1, channel_multiplier=0)

    g0_r = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    g0_i = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    g1_r = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    g1_i = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    nisa.nc_n_gather(dst=g0_r, data=X_real, indices=even_idx)
    nisa.nc_n_gather(dst=g0_i, data=X_imag, indices=even_idx)
    nisa.nc_n_gather(dst=g1_r, data=X_real, indices=odd_idx)
    nisa.nc_n_gather(dst=g1_i, data=X_imag, indices=odd_idx)

    # DFT
    d0_r = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    d0_i = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    d1_r = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    d1_i = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    _fft1d_matmul_isa(g0_r, g0_i, d0_r, d0_i, W_128_real, W_128_imag)
    _fft1d_matmul_isa(g1_r, g1_i, d1_r, d1_i, W_128_real, W_128_imag)

    # Butterfly
    tw_out_r = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    tw_out_i = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    tmp_ac = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    tmp_bd = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    tmp_ad = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    tmp_bc = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)

    Y_first_r = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    Y_first_i = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    Y_second_r = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)
    Y_second_i = nl.ndarray((TILE_H, N_half), dtype=nl.float32, buffer=nl.sbuf)

    _butterfly_pair(
        d0_r,
        d0_i,
        d1_r,
        d1_i,
        twiddle_real,
        twiddle_imag,
        Y_first_r,
        Y_first_i,
        Y_second_r,
        Y_second_i,
        tw_out_r,
        tw_out_i,
        tmp_ac,
        tmp_bd,
        tmp_ad,
        tmp_bc,
    )

    # Assemble output
    Y_combined_real = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.sbuf)
    Y_combined_imag = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.sbuf)
    nisa.tensor_copy(dst=Y_combined_real[0:TILE_H, 0:N_half], src=Y_first_r)
    nisa.tensor_copy(dst=Y_combined_imag[0:TILE_H, 0:N_half], src=Y_first_i)
    nisa.tensor_copy(dst=Y_combined_real[0:TILE_H, N_half:N], src=Y_second_r)
    nisa.tensor_copy(dst=Y_combined_imag[0:TILE_H, N_half:N], src=Y_second_i)

    nisa.dma_copy(dst=Y_real_hbm[0:TILE_H, 0:N], src=Y_combined_real)
    nisa.dma_copy(dst=Y_imag_hbm[0:TILE_H, 0:N], src=Y_combined_imag)
    return Y_real_hbm, Y_imag_hbm


# ---------------------------------------------------------------------------
# 512-pt NKI FFT kernel
# ---------------------------------------------------------------------------


@nki.jit
def _fft1d_512(
    X_real_hbm,
    X_imag_hbm,
    W_128_real_hbm,
    W_128_imag_hbm,
    tw256_real_hbm,
    tw256_imag_hbm,
    tw512_real_hbm,
    tw512_imag_hbm,
):
    """512-point FFT: 4 groups of 128, 2 butterfly levels."""
    N = 512

    Y_real_hbm = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.shared_hbm)
    Y_imag_hbm = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.shared_hbm)

    X_real = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.sbuf)
    X_imag = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.sbuf)
    nisa.dma_copy(dst=X_real, src=X_real_hbm[0:TILE_H, 0:N])
    nisa.dma_copy(dst=X_imag, src=X_imag_hbm[0:TILE_H, 0:N])

    W_128_real = nl.ndarray((BASE_N, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    W_128_imag = nl.ndarray((BASE_N, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    nisa.dma_copy(dst=W_128_real, src=W_128_real_hbm[0:BASE_N, 0:BASE_N])
    nisa.dma_copy(dst=W_128_imag, src=W_128_imag_hbm[0:BASE_N, 0:BASE_N])

    tw256_r = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    tw256_i = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    nisa.dma_copy(dst=tw256_r, src=tw256_real_hbm[0:TILE_H, 0:128])
    nisa.dma_copy(dst=tw256_i, src=tw256_imag_hbm[0:TILE_H, 0:128])

    tw512_r = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    tw512_i = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    nisa.dma_copy(dst=tw512_r, src=tw512_real_hbm[0:TILE_H, 0:256])
    nisa.dma_copy(dst=tw512_i, src=tw512_imag_hbm[0:TILE_H, 0:256])

    # Deinterleave into 4 groups
    idx0 = nl.ndarray((TILE_H, BASE_N), dtype=nl.uint32, buffer=nl.sbuf)
    idx1 = nl.ndarray((TILE_H, BASE_N), dtype=nl.uint32, buffer=nl.sbuf)
    idx2 = nl.ndarray((TILE_H, BASE_N), dtype=nl.uint32, buffer=nl.sbuf)
    idx3 = nl.ndarray((TILE_H, BASE_N), dtype=nl.uint32, buffer=nl.sbuf)
    nisa.iota(dst=idx0, pattern=[[4, 128]], offset=0, channel_multiplier=0)
    nisa.iota(dst=idx1, pattern=[[4, 128]], offset=1, channel_multiplier=0)
    nisa.iota(dst=idx2, pattern=[[4, 128]], offset=2, channel_multiplier=0)
    nisa.iota(dst=idx3, pattern=[[4, 128]], offset=3, channel_multiplier=0)

    g0_r = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    g0_i = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    g1_r = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    g1_i = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    g2_r = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    g2_i = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    g3_r = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    g3_i = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)

    nisa.nc_n_gather(dst=g0_r, data=X_real, indices=idx0)
    nisa.nc_n_gather(dst=g0_i, data=X_imag, indices=idx0)
    nisa.nc_n_gather(dst=g1_r, data=X_real, indices=idx1)
    nisa.nc_n_gather(dst=g1_i, data=X_imag, indices=idx1)
    nisa.nc_n_gather(dst=g2_r, data=X_real, indices=idx2)
    nisa.nc_n_gather(dst=g2_i, data=X_imag, indices=idx2)
    nisa.nc_n_gather(dst=g3_r, data=X_real, indices=idx3)
    nisa.nc_n_gather(dst=g3_i, data=X_imag, indices=idx3)

    # 128-pt DFT on each group
    d0_r = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    d0_i = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    d1_r = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    d1_i = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    d2_r = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    d2_i = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    d3_r = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)
    d3_i = nl.ndarray((TILE_H, BASE_N), dtype=nl.float32, buffer=nl.sbuf)

    _fft1d_matmul_isa(g0_r, g0_i, d0_r, d0_i, W_128_real, W_128_imag)
    _fft1d_matmul_isa(g1_r, g1_i, d1_r, d1_i, W_128_real, W_128_imag)
    _fft1d_matmul_isa(g2_r, g2_i, d2_r, d2_i, W_128_real, W_128_imag)
    _fft1d_matmul_isa(g3_r, g3_i, d3_r, d3_i, W_128_real, W_128_imag)

    # Level 1 butterfly (256-pt): pairs (g0,g2) and (g1,g3)
    tw_out_r = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    tw_out_i = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    tmp_ac = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    tmp_bd = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    tmp_ad = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    tmp_bc = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)

    b1_even_0_r = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    b1_even_0_i = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    b1_even_1_r = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    b1_even_1_i = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    _butterfly_pair(
        d0_r,
        d0_i,
        d2_r,
        d2_i,
        tw256_r,
        tw256_i,
        b1_even_0_r,
        b1_even_0_i,
        b1_even_1_r,
        b1_even_1_i,
        tw_out_r,
        tw_out_i,
        tmp_ac,
        tmp_bd,
        tmp_ad,
        tmp_bc,
    )

    b1_odd_0_r = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    b1_odd_0_i = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    b1_odd_1_r = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    b1_odd_1_i = nl.ndarray((TILE_H, 128), dtype=nl.float32, buffer=nl.sbuf)
    _butterfly_pair(
        d1_r,
        d1_i,
        d3_r,
        d3_i,
        tw256_r,
        tw256_i,
        b1_odd_0_r,
        b1_odd_0_i,
        b1_odd_1_r,
        b1_odd_1_i,
        tw_out_r,
        tw_out_i,
        tmp_ac,
        tmp_bd,
        tmp_ad,
        tmp_bc,
    )

    # Level 2 butterfly (512-pt)
    half_A_r = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    half_A_i = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    half_B_r = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    half_B_i = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)

    nisa.tensor_copy(dst=half_A_r[0:TILE_H, 0:128], src=b1_even_0_r)
    nisa.tensor_copy(dst=half_A_r[0:TILE_H, 128:256], src=b1_even_1_r)
    nisa.tensor_copy(dst=half_A_i[0:TILE_H, 0:128], src=b1_even_0_i)
    nisa.tensor_copy(dst=half_A_i[0:TILE_H, 128:256], src=b1_even_1_i)

    nisa.tensor_copy(dst=half_B_r[0:TILE_H, 0:128], src=b1_odd_0_r)
    nisa.tensor_copy(dst=half_B_r[0:TILE_H, 128:256], src=b1_odd_1_r)
    nisa.tensor_copy(dst=half_B_i[0:TILE_H, 0:128], src=b1_odd_0_i)
    nisa.tensor_copy(dst=half_B_i[0:TILE_H, 128:256], src=b1_odd_1_i)

    tw2_out_r = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    tw2_out_i = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    t2_ac = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    t2_bd = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    t2_ad = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    t2_bc = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)

    Y_first_r = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    Y_first_i = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    Y_second_r = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)
    Y_second_i = nl.ndarray((TILE_H, 256), dtype=nl.float32, buffer=nl.sbuf)

    _butterfly_pair(
        half_A_r,
        half_A_i,
        half_B_r,
        half_B_i,
        tw512_r,
        tw512_i,
        Y_first_r,
        Y_first_i,
        Y_second_r,
        Y_second_i,
        tw2_out_r,
        tw2_out_i,
        t2_ac,
        t2_bd,
        t2_ad,
        t2_bc,
    )

    # Assemble
    Y_combined_r = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.sbuf)
    Y_combined_i = nl.ndarray((TILE_H, N), dtype=nl.float32, buffer=nl.sbuf)
    nisa.tensor_copy(dst=Y_combined_r[0:TILE_H, 0:256], src=Y_first_r)
    nisa.tensor_copy(dst=Y_combined_i[0:TILE_H, 0:256], src=Y_first_i)
    nisa.tensor_copy(dst=Y_combined_r[0:TILE_H, 256:N], src=Y_second_r)
    nisa.tensor_copy(dst=Y_combined_i[0:TILE_H, 256:N], src=Y_second_i)

    nisa.dma_copy(dst=Y_real_hbm[0:TILE_H, 0:N], src=Y_combined_r)
    nisa.dma_copy(dst=Y_imag_hbm[0:TILE_H, 0:N], src=Y_combined_i)
    return Y_real_hbm, Y_imag_hbm


# =============================================================================
# PyTorch-level FFT interface (calls NKI kernels in tiles)
# =============================================================================


def _next_power_of_2(n):
    """Return the smallest power of 2 >= n."""
    if n <= 1:
        return 1
    return 1 << (n - 1).bit_length()


def _compute_dft_matrix(N=128):
    """Compute N-point DFT matrix split into real/imag."""
    k = np.arange(N, dtype=np.float32).reshape(N, 1)
    n = np.arange(N, dtype=np.float32).reshape(1, N)
    angles = -2.0 * np.pi * k * n / N
    return np.cos(angles), np.sin(angles)


def _compute_twiddle_factors(N, H=128):
    """Compute twiddle factors W_N^k for k=0..N/2-1, tiled to [H, N/2]."""
    k = np.arange(N // 2, dtype=np.float32)
    angles = -2.0 * np.pi * k / N
    return np.tile(np.cos(angles), (H, 1)), np.tile(np.sin(angles), (H, 1))


def nki_fft1d(x_real, x_imag, fft_size, device):
    """Perform 1D FFT using NKI kernels. Input shape: (H, W) where H <= 128.

    Pads W to fft_size (must be power of 2: 128, 256, or 512).
    Returns (Y_real, Y_imag) of shape (H, fft_size).
    """
    H, W = x_real.shape
    assert H <= 128, f"Tile height must be <= 128, got {H}"
    assert fft_size in (128, 256, 512), f"Only 128/256/512 supported, got {fft_size}"

    # Pad height to 128 and width to fft_size
    x_r = torch.zeros(128, fft_size, dtype=torch.float32, device=device)
    x_i = torch.zeros(128, fft_size, dtype=torch.float32, device=device)
    x_r[:H, :W] = x_real[:H, : min(W, fft_size)]
    x_i[:H, :W] = x_imag[:H, : min(W, fft_size)]

    # DFT matrix (always 128x128)
    W_r_np, W_i_np = _compute_dft_matrix(128)
    W_r = torch.tensor(W_r_np, dtype=torch.float32, device=device)
    W_i = torch.tensor(W_i_np, dtype=torch.float32, device=device)

    if fft_size == 128:
        y_r, y_i = _fft1d_128(x_r, x_i, W_r, W_i)
    elif fft_size == 256:
        tw_r_np, tw_i_np = _compute_twiddle_factors(256, 128)
        tw_r = torch.tensor(tw_r_np, dtype=torch.float32, device=device)
        tw_i = torch.tensor(tw_i_np, dtype=torch.float32, device=device)
        y_r, y_i = _fft1d_256(x_r, x_i, W_r, W_i, tw_r, tw_i)
    elif fft_size == 512:
        tw256_r_np, tw256_i_np = _compute_twiddle_factors(256, 128)
        tw512_r_np, tw512_i_np = _compute_twiddle_factors(512, 128)
        tw256_r = torch.tensor(tw256_r_np, dtype=torch.float32, device=device)
        tw256_i = torch.tensor(tw256_i_np, dtype=torch.float32, device=device)
        tw512_r = torch.tensor(tw512_r_np, dtype=torch.float32, device=device)
        tw512_i = torch.tensor(tw512_i_np, dtype=torch.float32, device=device)
        y_r, y_i = _fft1d_512(x_r, x_i, W_r, W_i, tw256_r, tw256_i, tw512_r, tw512_i)

    return y_r[:H, :fft_size], y_i[:H, :fft_size]


def _dft_matrix_torch(N, device):
    """Compute N-point DFT matrix as a torch tensor (for arbitrary N)."""
    k = torch.arange(N, dtype=torch.float32, device=device).unsqueeze(1)
    n = torch.arange(N, dtype=torch.float32, device=device).unsqueeze(0)
    angles = -2.0 * math.pi * k * n / N
    return torch.cos(angles), torch.sin(angles)


def _torch_fft1d(x_real, x_imag, device):
    """Pure PyTorch 1D FFT via DFT matrix multiply. For sizes not supported by NKI.

    Input: (H, N) real and imag tensors.
    Output: (H, N) real and imag tensors.
    """
    H, N = x_real.shape
    W_r, W_i = _dft_matrix_torch(N, device)  # (N, N)
    # Y = X @ W^T (complex)
    # Y_real = X_real @ W_r^T - X_imag @ W_i^T
    # Y_imag = X_real @ W_i^T + X_imag @ W_r^T
    y_real = x_real @ W_r.t() - x_imag @ W_i.t()
    y_imag = x_real @ W_i.t() + x_imag @ W_r.t()
    return y_real, y_imag


def _fft1d_dispatch(x_real, x_imag, N, device):
    """Dispatch to NKI kernel for supported sizes, else use PyTorch DFT.

    Performs an N-point FFT on each row. Input must have exactly N columns.
    """
    if N in (128, 256, 512) and x_real.shape[0] <= 128:
        return nki_fft1d(x_real, x_imag, N, device)
    else:
        # Use PyTorch matrix-multiply DFT for unsupported sizes
        return _torch_fft1d(x_real, x_imag, device)


def nki_fft2d_real(hidden_states, device):
    """Perform 2D FFT on hidden_states (B, S, D) and return real part.

    Equivalent to torch.fft.fftn(hidden_states, dim=(1, 2)).real

    Decomposes into:
    1. 1D FFT along dim=2 (hidden_dim, size D) for each row
    2. 1D FFT along dim=1 (seq_dim, size S) for each column
    3. Take real part

    Uses NKI kernels for sizes 128/256/512; falls back to PyTorch DFT
    matrix multiply for other sizes.
    """
    B, S, D = hidden_states.shape

    # Process batch dimension
    results = []
    for b in range(B):
        x = hidden_states[b].float()  # (S, D)

        # --- Pass 1: FFT along hidden_dim (dim=1, size D) ---
        # Process in tiles of 128 rows for NKI compatibility
        pass1_real = torch.zeros(S, D, dtype=torch.float32, device=device)
        pass1_imag = torch.zeros(S, D, dtype=torch.float32, device=device)

        tile_size = 128 if D in (128, 256, 512) else S  # Only tile for NKI sizes

        for tile_start in range(0, S, 128):
            tile_end = min(tile_start + 128, S)
            tile_h = tile_end - tile_start
            tile_data = x[tile_start:tile_end]  # (tile_h, D)
            tile_imag = torch.zeros_like(tile_data)

            yr, yi = _fft1d_dispatch(tile_data, tile_imag, D, device)
            pass1_real[tile_start:tile_end] = yr[:tile_h, :D]
            pass1_imag[tile_start:tile_end] = yi[:tile_h, :D]

        # --- Pass 2: FFT along seq_dim (dim=0, size S) ---
        # Transpose: (S, D) -> (D, S), FFT along S, transpose back
        p1_r_t = pass1_real.t().contiguous()  # (D, S)
        p1_i_t = pass1_imag.t().contiguous()  # (D, S)

        pass2_real = torch.zeros(D, S, dtype=torch.float32, device=device)
        pass2_imag = torch.zeros(D, S, dtype=torch.float32, device=device)

        for tile_start in range(0, D, 128):
            tile_end = min(tile_start + 128, D)
            tile_h = tile_end - tile_start

            tile_r = p1_r_t[tile_start:tile_end]  # (tile_h, S)
            tile_i = p1_i_t[tile_start:tile_end]

            yr, yi = _fft1d_dispatch(tile_r, tile_i, S, device)
            pass2_real[tile_start:tile_end] = yr[:tile_h, :S]
            pass2_imag[tile_start:tile_end] = yi[:tile_h, :S]

        # Transpose back: (D, S) -> (S, D)
        final_real = pass2_real.t()  # (S, D) - real part only
        results.append(final_real)

    return torch.stack(results, dim=0)  # (B, S, D)


# =============================================================================
# Layout Class
# =============================================================================


class NeuronFNetFourierLayout(nn.Module):
    """Layout for FNetBasicFourierTransform. No learnable weights needed."""

    conversion_mapping = []

    def __init__(self, config, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        # Store config for TPU optimization path compatibility
        self._use_tpu_fourier = getattr(config, "use_tpu_fourier_optimizations", False)

    def forward(self, hidden_states):
        pass


# =============================================================================
# Forward Class
# =============================================================================


class NeuronFNetFourierForward(nn.Module):
    """NKI-accelerated 2D Fourier Transform for FNet.

    Replaces torch.fft.fftn(x, dim=(1,2)).real with NKI kernels that perform
    the FFT entirely on NeuronCore hardware using the Tensor Engine for the
    128-point DFT base case.
    """

    def forward(self, hidden_states):
        device = hidden_states.device

        # NKI 2D FFT (real part only)
        output = nki_fft2d_real(hidden_states, device)

        # Cast back to model dtype
        output = output.to(hidden_states.dtype)
        return (output,)


# =============================================================================
# Layer Registry
# =============================================================================


class layers:
    NeuronFNetFourierForward = NeuronFNetFourierForward