Example use of the pyvkfft.fft interface

Using this interface, the explicit VkFFTApp (a.k.a. the fft ‘plan’) creation is not necessary, they are automatically created and cached for future re-use.

Also, the appropriate backend for a pycuda/pyopencl or cupy array is automatically selected.

[1]:
# Use the following to install pyvkfft on google colab
if False:
    # Install pyvkfft & dependencies only once using pip
    import os
    if not os.path.exists('dev/pyvkfft'):
      !mkdir dev
      !cd dev && git clone --recursive https://github.com/vincefn/pyvkfft.git
      !pip install pycuda
      # !pip install cupy
      !cd dev/pyvkfft && pip install .
      # scipy, matplotlib not required for pyvkfft, but for tests
      !pip install scipy matplotlib
[2]:
from pyvkfft.fft import fftn, ifftn, rfftn, irfftn

import numpy as np
try:
    from scipy.datasets import ascent
except ImportError:
    from scipy.misc import ascent
from numpy.fft import fftshift
%matplotlib notebook
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
[3]:
import sys
sys.path
[3]:
['/home/esrf/favre/dev/pyvkfft/examples',
 '/home/esrf/favre/miniconda3/envs/pynx-py311-cu11.7/lib/python311.zip',
 '/home/esrf/favre/miniconda3/envs/pynx-py311-cu11.7/lib/python3.11',
 '/home/esrf/favre/miniconda3/envs/pynx-py311-cu11.7/lib/python3.11/lib-dynload',
 '',
 '/home/esrf/favre/miniconda3/envs/pynx-py311-cu11.7/lib/python3.11/site-packages']

Test functions

These only use the fftn, ifftn, rfftn and irfftn functions, which will automatically detect the type of GPU array, and cache the generated VkFFTApp.

[4]:
def do_fft_and_plot(d):
    plt.figure(figsize=(9,3))

    plt.subplot(131)
    plt.imshow(abs(d.get()))
    plt.colorbar()

    d = fftn(d,norm="ortho")  # Use fftn(d,d) for an inplace transform

    plt.subplot(132)
    plt.imshow(fftshift(abs(d.get())), norm=LogNorm())
    plt.colorbar()

    d = ifftn(d,norm="ortho")  # Use fftn(d,d) for an inplace transform

    plt.subplot(133)
    plt.imshow(abs(d.get()))
    plt.colorbar()
    plt.tight_layout()

def do_rfft_and_plot(dr, dc=None):
    # if dc is None, the transform is out-of-place and the destination
    # array is allocated on-the-fly
    plt.figure(figsize=(9,3))

    plt.subplot(131)
    plt.imshow(abs(d.get()))
    plt.colorbar()

    dc = rfftn(dr, dc,norm="ortho")

    plt.subplot(132)
    plt.imshow(fftshift(abs(dc.get()), axes=[0]), norm=LogNorm())
    plt.colorbar()

    dr = irfftn(dc, dr,norm="ortho")

    plt.subplot(133)
    plt.imshow(abs(dr.get()))
    plt.colorbar()
    plt.tight_layout()

PyCUDA

[5]:
try:
    import pycuda.autoinit
    import pycuda.gpuarray as cua
    has_pycuda = True
except:
    has_pycuda = False

if has_pycuda:
    print("C2C transform")
    d = cua.to_gpu(ascent().astype(np.complex64))
    do_fft_and_plot(d)

    print("R2C transform, out-of-place in auto-allocated array")
    d = cua.to_gpu(ascent().astype(np.float32))
    do_rfft_and_plot(d)

    print("R2C transform, out-of-place")
    dr = cua.to_gpu(ascent().astype(np.float32))
    sh = (dr.shape[0], dr.shape[1]//2+1)
    dc = cua.empty(sh, dtype=np.complex64)
    do_rfft_and_plot(dr, dc)

    print("R2C transform, inplace")
    dr = cua.to_gpu(ascent().astype(np.float32))
    do_rfft_and_plot(dr, dr)
else:
    print("PyCUDA is not available")
C2C transform
R2C transform, out-of-place in auto-allocated array
R2C transform, out-of-place
R2C transform, inplace

CuPy

[6]:
try:
    import cupy as cp
    has_cupy = True
except:
    has_cupy = False

if has_cupy:
    print("C2C transform")
    d = cp.array(ascent().astype(np.complex64))
    do_fft_and_plot(d)

    print("R2C transform, out-of-place in auto-allocated array")
    d = cp.array(ascent().astype(np.float32))
    do_rfft_and_plot(d)

    print("R2C transform, out-of-place")
    dr = cp.array(ascent().astype(np.float32))
    sh = (dr.shape[0], dr.shape[1]//2+1)
    dc = cp.empty(sh, dtype=np.complex64)
    do_rfft_and_plot(dr, dc)

    print("R2C transform, inplace")
    dr = cp.array(ascent().astype(np.float32))
    do_rfft_and_plot(dr, dr)
else:
    print("CuPy is not available")
C2C transform
R2C transform, out-of-place in auto-allocated array
R2C transform, out-of-place
R2C transform, inplace

PyOpenCL

[8]:
try:
    import pyopencl as cl
    import pyopencl.array as cla
    import os

    # Create some context on the first available GPU
    if 'PYOPENCL_CTX' in os.environ:
        ctx = cl.create_some_context()
    else:
        ctx = None
        # Find the first OpenCL GPU available and use it, unless
        for p in cl.get_platforms():
            for d in p.get_devices():
                if d.type & cl.device_type.GPU == 0:
                    continue
                print("Selected device: ", d.name)
                ctx = cl.Context(devices=(d,))
                break
            if ctx is not None:
                break
    cq = cl.CommandQueue(ctx)

    has_pyopencl = True
except:
    has_pyopencl = False

if has_pyopencl:
    print("C2C transform")
    d = cla.to_device(cq, ascent().astype(np.complex64))
    do_fft_and_plot(d)

    print("R2C transform, out-of-place in auto-allocated array")
    d = cla.to_device(cq, ascent().astype(np.float32))
    do_rfft_and_plot(d)

    print("R2C transform, out-of-place")
    dr = cla.to_device(cq, ascent().astype(np.float32))
    sh = (dr.shape[0], dr.shape[1]//2+1)
    dc = cla.empty(cq, sh, dtype=np.complex64)
    do_rfft_and_plot(dr, dc)

    print("R2C transform, inplace")
    dr = cla.to_device(cq, ascent().astype(np.float32))
    do_rfft_and_plot(dr, dr)
else:
    print("PyOpenCL is not available")
Selected device:  NVIDIA A40
C2C transform
R2C transform, out-of-place in auto-allocated array
R2C transform, out-of-place
R2C transform, inplace
[ ]: