warp.tile_fft#

warp.tile_fft(
inout: Tile[Vector[Float, Literal[2]], tuple[int, ...]],
) → None#
  • Kernel
  • Differentiable

Compute the forward FFT along the last dimension of an N-D tile, in place.

All leading dimensions are treated as independent batches and the last dimension is the transform size N. The tile must have at least two dimensions and register storage. inout is overwritten with its transform.

Each element is a vec2 interpreted as a complex number, with component 0 the real part and component 1 the imaginary part. The transform is the standard unnormalized DFT, X[k] = sum_n x[n] * exp(-2*pi*i*k*n/N), with output bins in the usual order k = 0 ... N-1, matching numpy.fft.fft(). Because no scaling is applied, a forward transform followed by an inverse transform multiplies the data by N; normalization is left to the caller.

Supported dtypes are vec2f and vec2d.

This is a cooperative operation, so every thread in the block must call it.

Backward propagation applies an inverse FFT along the last dimension to the incoming gradient, without dividing by N.

A MathDx-enabled build is not required, but the supported transform sizes depend on the backend:

  • CPU: any power-of-two size; non-power-of-two sizes go through an O(N^2) DFT and support at most 4096 elements per transform.

  • GPU with libmathdx (cuFFTDx), unless disabled by warp.config.enable_mathdx_fft or the enable_mathdx_fft module option: supported N and block_dim combinations depend on the build. block_dim must be at most N/2.

  • GPU without libmathdx, or with cuFFTDx disabled: power-of-two sizes only, and N must be divisible by block_dim. Compiles faster but runs slower than the cuFFTDx path.

Unsupported transform sizes or block dimensions raise an exception during kernel compilation, usually on the first launch.

Parameters:

inout – The tile to transform in place.

Example

@wp.kernel
def spectrum(x: wp.array2d[wp.vec2f]):
    x_tile = wp.tile_load(x, shape=(2, 4))
    wp.tile_fft(x_tile)
    wp.tile_store(x, x_tile)

# two batches of four real samples; the FFT runs along the last dimension
signal = np.zeros((2, 4, 2), dtype=np.float32)
signal[0, :, 0] = [1.0, 2.0, 3.0, 4.0]
signal[1, :, 0] = [4.0, 3.0, 2.0, 1.0]
x = wp.array(signal, dtype=wp.vec2f)

# each thread owns two samples from every batch
wp.launch_tiled(spectrum, dim=1, inputs=[x], block_dim=2)

print(np.round(x.numpy(), 3))

[[[10.  0.]
  [-2.  2.]
  [-2.  0.]
  [-2. -2.]]

 [[10.  0.]
  [ 2. -2.]
  [ 2.  0.]
  [ 2.  2.]]]