[REQ] Support N-D tiles in `wp.tile_fft()` and `wp.tile_ifft()`

[ns5678]

Description

Request to lift the 2D-tile-only restriction so wp.tile_fft() works on tiles with 3+ dimensions, allowing FFT along the last axis regardless of tile dimensionality.

Right now, the code below would error out with the message tile_fft() requires at least 2 elements per thread, but got ept=0 (fft_size=1, block_dim=32). Reduce block_dim to at most 0 for this FFT size. This happens because in def tile_fft_generic_lto_dispatch_func() inside warp/_src/builtins.py, batch, size = inout.type.shape[0], inout.type.shape[1] allowing FFT for only 2-D tiles.

import numpy as np

import warp as wp

N = 128
BLOCK_DIM = 32

@wp.kernel(module="fft_3d_kernel")
def fft_3d_kernel(gx: wp.array3d(dtype=wp.vec2f), gy: wp.array3d(dtype=wp.vec2f)):
    x = wp.tile_load(gx, shape=(1, 1, N))
    wp.tile_fft(x)
    wp.tile_store(gy, x)


if __name__ == "__main__":
    rng = np.random.default_rng(42)
    X = rng.random((N, N, 2 * N), dtype=np.float32)

    X_wp = wp.array(X, dtype=wp.vec2f, ndim=3)
    Y_wp = wp.zeros_like(X_wp)

    wp.launch_tiled(
        fft_3d_kernel,
        dim=[N, N],
        inputs=[X_wp, Y_wp],
        block_dim=BLOCK_DIM,
    )

Context

This feature would be helpful in efficiently extending the warp/examples/core/example_fft_poisson_navier_stokes_2d.py to 3-D setup. In general, this could also help with scenarios where one is dealing with data in >2-D dimensions and need FFT/IFFT operations.

[shi-eric]

Added in ff3d8d9