Metadata-Version: 2.4
Name: trnfft
Version: 0.27.0
Summary: FFT and complex-valued tensor operations for AWS Trainium via NKI
Author-email: Scott Friedman <scttfrdmn@gmail.com>
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/trnsci/trnfft
Project-URL: Documentation, https://trnsci.dev/trnfft/
Project-URL: Repository, https://github.com/trnsci/trnfft
Project-URL: Issues, https://github.com/trnsci/trnfft/issues
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1
Requires-Dist: numpy>=1.24
Provides-Extra: neuron
Requires-Dist: nki>=0.3.0; extra == "neuron"
Requires-Dist: torch-neuronx>=2.9; extra == "neuron"
Provides-Extra: torchneuron
Requires-Dist: torch>=2.14; extra == "torchneuron"
Requires-Dist: torch-neuronx>=2.14; extra == "torchneuron"
Requires-Dist: nki>=0.6; extra == "torchneuron"
Provides-Extra: dev
Requires-Dist: pytest<9,>=7.0; extra == "dev"
Requires-Dist: pytest-benchmark>=4.0; extra == "dev"
Requires-Dist: pytest-cov>=4.0; extra == "dev"
Requires-Dist: ruff>=0.8; extra == "dev"
Requires-Dist: pre-commit>=3.5; extra == "dev"
Requires-Dist: scipy>=1.11; extra == "dev"
Requires-Dist: mkdocs>=1.5; extra == "dev"
Requires-Dist: mkdocs-material>=9.0; extra == "dev"
Dynamic: license-file

# trnfft

[![CI](https://github.com/trnsci/trnfft/actions/workflows/ci.yml/badge.svg)](https://github.com/trnsci/trnfft/actions/workflows/ci.yml)
[![codecov](https://codecov.io/gh/trnsci/trnfft/branch/main/graph/badge.svg)](https://codecov.io/gh/trnsci/trnfft)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
[![PyPI](https://img.shields.io/pypi/v/trnfft)](https://pypi.org/project/trnfft/)
[![Python](https://img.shields.io/pypi/pyversions/trnfft)](https://pypi.org/project/trnfft/)
[![License](https://img.shields.io/github/license/trnsci/trnfft)](LICENSE)
[![Docs](https://img.shields.io/badge/docs-mkdocs-blue)](https://trnsci.github.io/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](https://github.com/scttfrdmn/neuron-complex-ops). Part of the trnsci scientific computing suite ([github.com/trnsci](https://github.com/trnsci)).

## Current phase

trnfft follows the [trnsci 5-phase roadmap](https://trnsci.dev/roadmap/). Active work is tracked in phase-labeled GitHub issues:

- **[Phase 1 — correctness](https://github.com/trnsci/trnfft/issues/51)**: **complete as of v0.8.0** (butterfly + complex GEMM kernels hardware-validated on trn1.2xlarge, 70/70 benchmark cases passing).
- **[Phase 2 — precision](https://github.com/trnsci/trnfft/issues/52)** (active): Kahan / Neumaier compensated summation. Partial delivery in v0.11.0 (precision-modes API + compensated butterfly). On-silicon kahan characterization (#58) and Thread C (Tensor-engine Stockham) are the remaining items.
- **[Phase 3 — perf](https://github.com/trnsci/trnfft/issues/53)**: plan reuse, streaming large FFTs, NEFF cache reuse.
- **[Phase 4 — multi-chip](https://github.com/trnsci/trnfft/issues/54)**: multi-chip large-N FFT (N > 2²⁰).
- **[Phase 5 — generation](https://github.com/trnsci/trnfft/issues/55)**: trn2 larger-SBUF butterfly path.

Suite-wide tracker: [trnsci/trnsci#1](https://github.com/trnsci/trnsci/issues/1).

## 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

```bash
pip install trnfft

# Trn2/Trn3 with TorchNeuron (native PyTorch, Neuron SDK 2.33+)
pip install "trnfft[torchneuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com

# Trn1 / PyTorch-XLA lineage
pip install "trnfft[neuron]" --extra-index-url https://pip.repos.neuron.amazonaws.com
```

Kernels dispatch on whichever Neuron PyTorch lineage is installed:

- **TorchNeuron (native)**: Trn2/Trn3, Neuron SDK 2.33+, Python ≥ 3.12,
  `torch` 2.14, `torch-neuronx` 2.14.*. Each kernel is a PyTorch custom op
  (`torch.ops.trnfft.*`) on the `neuron` device; no `torch_xla`. Gradients
  flow through `torch.library.register_autograd`.
- **PyTorch/XLA**: Trn1 (and any host on the `torch-neuronx` 2.9 lineage).
  Kernels dispatch through `torch_xla`, unchanged.

trnfft picks the native path when the native `torch_neuronx` and a `neuron`
device are present, and the XLA path when `torch_xla` is installed. Set
`TRNSCI_NKI_DISPATCH=native`, `xla` or `cpu` (shared by all trnsci libraries)
to force a path. With an editable install (`pip install -e`), set
`NKI_ENABLE_TRACE_CACHE=0`: the NKI kernel cache only versions installed
wheels and can serve stale kernels. See
[installation](docs/installation.md).

### Known platform issue: NKI trace cache and kernel factories

`import torch_neuronx` turns on the NKI trace cache, and the cache key ignores
constants captured in a kernel's closure: two `@nki.jit` kernels built by one
factory, differing only in a captured value, share one trace, and the second
silently returns the first's results
([aws-neuron/aws-neuron-sdk#1454](https://github.com/aws-neuron/aws-neuron-sdk/issues/1454)).
trnfft's kernels do not rely on closure constants (or on module globals that
can differ between processes): every value that specialises a kernel is a
kernel argument, and `tests/test_trace_cache_safety.py` enforces this. If you
write your own NKI kernels with a factory, set `NKI_ENABLE_TRACE_CACHE=0`
before importing `torch_neuronx`.

## Usage

```python
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 (torch.stft's arguments; window=None means Hann here, ones in torch.stft)
waveform = torch.randn(16000)
S = trnfft.stft(waveform, n_fft=512, hop_length=256, window=torch.hann_window(512))
```

## Complex Neural Network Layers

```python
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) on the PyTorch backend. On the NKI path, FP32: a single DFT-GEMM for N < 512, the batch-fused four-step (two GEMMs) up to N = 4096, and the six-step for N ≥ 8192 (the default since v0.25.0; `TRNFFT_SIX_STEP=0` restores the Stockham / butterfly path)
- **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 on Trn2 (trn2.3xlarge) against **Neuron SDK 2.33** with native TorchNeuron, and on Trn1 against **Neuron SDK 2.29+** on the **Deep Learning AMI Neuron PyTorch 2.9 (Ubuntu 24.04)** AMI (XLA lineage). See [docs/installation.md](https://trnsci.github.io/trnfft/installation/#hardware-compatibility) for the full compatibility matrix.

## Benchmarks

NKI vs PyTorch on the same Trainium instance — see the [benchmarks page](https://trnsci.github.io/trnfft/benchmarks/) for the latest numbers.

## Status

**v0.26.0** — validated on Trn2 (Neuron SDK 2.33, native TorchNeuron). Correctness releases, upgrade recommended: 0.25.0 fixed GEMM tail tiles that ≤ 0.24.0 dropped silently (small-N DFT-GEMM, bf16 precision modes, `nn.ComplexLinear`, `complex_gemm`); 0.26.0 fixed `irfft` / `hfft` / `irfft2` / `irfftn` using every bin instead of the first `n//2+1` (#105) and float64 FFTs on the NKI path returning inf/NaN (#109). See the [CHANGELOG](CHANGELOG.md).

**API coverage** (24 `torch.fft` functions):
`fft`, `ifft`, `rfft`, `irfft`, `hfft`, `ihfft`, `fft2`, `ifft2`, `rfft2`, `irfft2`, `hfft2`, `ihfft2`, `fftn`, `ifftn`, `rfftn`, `irfftn`, `hfftn`, `ihfftn`, `fftshift`, `ifftshift`, `fftfreq`, `rfftfreq`, `stft`, `istft`. Every transform takes `torch.fft`'s `n` / `s`, `dim` and `norm` arguments with the same meaning (#123).

**Roadmap**
- NKI `ComplexConv1d` / `ComplexModReLU` kernels (today both fall back to PyTorch on NKI)
- BF16 beyond the N ≤ 256 DFT-GEMM precision modes (`bf16`, `bf16_refined`, `ozaki`, `ozaki_hq`); FP16
- Multi-NeuronCore parallelism: `set_multicore` splits host batches across cores, but
  on Trn2 the host link limits it (#93); large N runs best on one core, on the
  six-step path that is the default from N=8192 (#94; `TRNFFT_SIX_STEP=0` turns it off)
- SBUF-resident dispatch to reduce small-op overhead

## Related projects in the trnsci suite

All six siblings are on PyPI, along with the umbrella meta-package:

| Project | What | Latest |
|---------|------|-------:|
| [trnsci](https://github.com/trnsci/trnsci) | Umbrella meta-package pulling the whole suite | v0.1.0 |
| [trnblas](https://github.com/trnsci/trnblas) | BLAS Level 1–3 for Trainium | v0.9.1 |
| [trnrand](https://github.com/trnsci/trnrand) | Philox / Sobol / Halton random number generation | v0.12.1 |
| [trnsolver](https://github.com/trnsci/trnsolver) | Linear solvers (CG, GMRES) and eigendecomposition | v0.20.0 |
| [trnsparse](https://github.com/trnsci/trnsparse) | Sparse matrix operations | v0.12.1 |
| [trntensor](https://github.com/trnsci/trntensor) | Tensor contractions (einsum, TT/Tucker decompositions) | v0.22.0 |
| [neuron-complex-ops](https://github.com/scttfrdmn/neuron-complex-ops) | Original proof-of-concept, folded into trnfft | archived |

## License

Apache 2.0 — Copyright 2026 Scott Friedman

## Disclaimer

trnsci is an **independent open-source project**. It is not sponsored by, endorsed by, or affiliated with Amazon.com, Inc., Amazon Web Services, Inc., or Annapurna Labs Ltd.

"AWS", "Amazon", "Trainium", "Inferentia", "NeuronCore", "Neuron SDK", and related identifiers are trademarks of their respective owners and are used here solely for descriptive and interoperability purposes. Use does not imply endorsement, partnership, or any other relationship.

All work, opinions, analyses, benchmark results, architectural commentary, and editorial judgments in this repository and on [trnsci.dev](https://trnsci.dev) are those of the project's contributors. They do not represent the views, positions, or commitments of Amazon, AWS, or Annapurna Labs.

Feedback directed at the Neuron SDK or Trainium hardware is good-faith ecosystem commentary from independent users. It is not privileged information, is not pre-reviewed by AWS, and should not be read as authoritative about product roadmap, behavior, or quality.

For official AWS guidance, see [aws-neuron documentation](https://awsdocs-neuron.readthedocs-hosted.com/) and the [AWS Trainium product page](https://aws.amazon.com/ai/machine-learning/trainium/).

## Acknowledgments

Built on insights from:
- [tcFFT](https://github.com/tcFFT) — Tensor Core FFT research
- [FFTW](https://www.fftw.org/) — Plan-based FFT architecture
- [AWS NKI documentation](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/nki/index.html)
