Add 1D Fast Fourier Transform FFT NKI kernel implementation

[jimburtoft]

Description of changes:

• Hardware-accelerated 1D FFT using Tensor Engine
• Supports arbitrary heights and power-of-2 widths (128-4096)
• Radix-2 Cooley-Tukey algorithm with 128-point DFT base case
• High accuracy (< 0.003% error vs NumPy)
• Comprehensive unit tests with performance benchmarks"

Testing:

Just covered by function specific unit testing.

Pull Request Checklist

• [x ] I have filled in all the required field in the template
• [ x] I have tested locally that all the tests pass
• [x ] By submitting this pull request, I confirm that my contribution is made under the terms of the MIT-0 license.

[mrkcath-aws]

Thank you for the submission!

Feedback: the kernel imports neuronxcc.nki, which is the deprecated NKI namespace. The current API is the top-level nki package, and a few of the specific calls here changed in the move (nl.mgrid and mask= on nl.load/nl.store were removed; nisa.nc_matmul/nc_transpose are now dest-style, nc_matmul(dst, stationary, moving)).

I attempted the migration and ran it on a Trainium2 instance against a numpy.fft.fft reference across your full test matrix and it holds up well. All under your stated 0.003%. The core is unchanged, with the same radix-2 Cooley-Tukey with the 128-point Tensor-Engine DFT base, same arbitrary-height (1–128) and power-of-2-width (128–4096) support, same accuracy.

The migration notes, in case they're useful:

• The DFT matrix is symmetric (W[k,n] = W[n,k]), so X @ Wᵀ == X @ W and the migrated version transposes only the inputs, not W.
• Elementwise ops can't take two PSUM operands, so the four sub-matmul results are evacuated to SBUF before the real/imag combine.
• One feature was dropped in the port: the height-masking. nl.mgrid and masked load/store don't exist in the new API. This does not effect results.

Happy to share the migrated file if it's helpful as a starting point.

[jimburtoft]

updated kernel available at https://huggingface.co/kernels/jburtoft/fnet-fast-fourier-transform-neuron-kernels