trnfft 0.6.0

FFT and complex-valued tensor operations for AWS Trainium via NKI

trnfft

FFT and complex-valued tensor operations for AWS Trainium via NKI.

Trainium has no native complex number support and ships no FFT library. trnfft fills that gap with split real/imaginary representation, complex neural network layers, and NKI kernels optimized for the NeuronCore architecture.

Incorporates neuron-complex-ops. Part of the trn-* scientific computing suite by Playground Logic.

Why

NVIDIA has cuFFT, cuBLAS, and native complex64. Trainium has none of these. Every signal processing, speech enhancement, physics simulation, and spectral method workload on Trainium currently falls back to CPU or requires hand-rolling complex arithmetic. trnfft fixes this.

Install

pip install trnfft

# With Neuron hardware support
pip install trnfft[neuron]

Usage

import torch
import trnfft

# Drop-in replacement for torch.fft
signal = torch.randn(1024)
X = trnfft.fft(signal)
recovered = trnfft.ifft(X)

# Real-valued FFT
X = trnfft.rfft(signal)

# 2D FFT
image = torch.randn(256, 256)
F = trnfft.fft2(image)

# STFT (matches torch.stft signature)
waveform = torch.randn(16000)
S = trnfft.stft(waveform, n_fft=512, hop_length=256)

Complex Neural Network Layers

from trnfft import ComplexTensor
from trnfft.nn import ComplexLinear, ComplexConv1d, ComplexModReLU

# Build complex-valued models for speech/audio/physics
x = ComplexTensor(real_part, imag_part)
layer = ComplexLinear(256, 128)
y = layer(x)

Architecture

+--------------------------------------------+
|            User Code / Model               |
+--------------------------------------------+
|         trnfft.api (torch.fft API)         |
|   fft()  ifft()  rfft()  stft()  fft2()   |
+--------------------------------------------+
|   trnfft.fft_core     |  trnfft.nn        |
|   Cooley-Tukey         |  ComplexLinear    |
|   Bluestein            |  ComplexConv1d    |
|   Plan caching         |  ComplexModReLU   |
+------------------------+-------------------+
|       trnfft.nki.dispatch                  |
|   "auto" | "pytorch" | "nki"              |
+--------------------------------------------+
|  PyTorch ops     |  NKI kernels           |
|  (any device)    |  (Trainium only)       |
|  torch.matmul    |  nisa.nc_matmul        |
|  element-wise    |  Tensor Engine         |
|                  |  Vector Engine          |
|                  |  SBUF ↔ PSUM pipeline  |
+------------------+------------------------+

How It Works

No complex dtype? Trainium's NKI doesn't support complex64/complex128. ComplexTensor stores complex values as paired real tensors and decomposes complex arithmetic into real-valued operations.

FFT → butterflies → matmul. Each Cooley-Tukey butterfly stage performs complex-multiply-and-add across all groups simultaneously. On NKI, the complex multiply maps to the Tensor Engine (systolic array).

Algorithms:

• Power-of-2: Cooley-Tukey radix-2 (iterative, decimation-in-time)
• Arbitrary sizes: Bluestein's chirp-z transform (pads to power-of-2)

NKI complex GEMM uses stationary tile reuse (2 SBUF loads instead of 8) and PSUM accumulation, overlapping Vector Engine negation with Tensor Engine matmul.

Hardware compatibility

NKI kernels are validated against Neuron SDK 2.24+ on the Deep Learning AMI Neuron PyTorch 2.9 (Ubuntu 24.04) AMI (20260410 or later). See docs/installation.md for the full compatibility matrix.

Status

v0.1.0 — CPU fallback works, NKI kernels scaffolded for on-hardware validation.

• ComplexTensor with full arithmetic
• Complex matmul (4 real matmuls)
• 1D FFT/IFFT (power-of-2, Cooley-Tukey)
• Bluestein (arbitrary sizes)
• rfft/irfft
• 2D FFT
• STFT
• Complex NN layers (Linear, Conv1d, BatchNorm, ModReLU)
• NKI dispatch layer (auto/pytorch/nki)
• Plan caching
• NKI butterfly kernel validation on trn1/trn2
• NKI GEMM kernel validation
• Multi-NeuronCore parallelism
• Benchmarks vs cuFFT
• Inverse STFT
• N-D FFT

Related Projects

Project	What
neuron-complex-ops	Original proof-of-concept (now folded into this library)
trnblas	BLAS for Trainium (Level 1-3, DF-MP2 use case)
trnsolver (planned)	Linear solvers and eigendecomposition for Trainium

License

Apache 2.0 — Playground Logic LLC

Acknowledgments

Built on insights from:

• tcFFT — Tensor Core FFT research
• FFTW — Plan-based FFT architecture
• AWS NKI documentation