pyvkfft
pyvkfft - python interface to the CUDA and OpenCL backends of VkFFT (Vulkan Fast Fourier Transform library)
VkFFT is a GPU-accelerated Fast Fourier Transform library for Vulkan/CUDA/HIP/OpenCL.
pyvkfft offers a simple python interface to the CUDA and OpenCL backends of VkFFT, compatible with pyCUDA, CuPy and pyOpenCL.
The documentation can be found at https://pyvkfft.readthedocs.io
See the:
Installation
Requirements: pyvkfft
only requires numpy
, psutil
, and at least one
GPU computing package among pyopencl
, pycuda
and cupy
.
You can install pyvkfft
either using PyPI:
pip install pyvkfft
or conda (using the conda-forge channel):
conda config --add channels conda-forge
conda install pyvkfft
When using cuda, you can also specify the nvrtc library version
using the cuda-version
package:
# Example also with pyopencl and cupy
conda install pyvkfft pyopencl cupy cuda-version=12
Driver/toolkit requirements
pyvkfft
needs a working GPU computing environment, including drivers,
compiling tools and libraries.
For OpenCL, the Installable Client Driver (ICD) should be available - it is normally
provided automatically when installing GPU toolkits (cuda, orcm,...). When using conda,
it may be necessary to install ocl-icd-system
(under Linux) or ocl-icd-wrapper-apple
.
This is done automatically if pyvkfft
is installed using conda.
For CUDA, in addition to the driver, the cuda toolkit (including nvcc
and the
real-time-compute (rtc) libraries) must also be installed. The compiler and toolkit
can be installed using conda
if they are not available system-wide.
When installing using pip (needs compilation), the path to nvcc
(or nvcc.exe
) will be automatically searched, first using the CUDA_PATH
or CUDA_HOME
environment variables, or then in the PATH
.
If nvcc
is not found, only support for OpenCL will be compiled.
Windows installation (cuda)
Windows installation can be tricky. Here are some hints for windows 10 with an nvidia card - YMMV:
first install Visual studio (tested with 2019) with C++ tools and windows SDK
install conda using mambaforge which has the advantage of including
conda-forge
as a default channel, andmamba
for faster installations
Then you can either install all packages using conda:
mamba create -n myenv python=3.11 numpy scipy matplotlib ipython psutil pyopencl cupy pycuda nvidia::cuda pyvkfft
or install all packages except pyvkfft using conda:
mamba create -n myenv python=3.11 numpy scipy matplotlib ipython psutil pyopencl cupy pycuda nvidia::cuda
conda activate myenv
pip install pyvkfft
or install using conda for base packages including nvidia::cuda, then pypi for pyopencl, pycuda, cupy and pyvkfft
mamba create -n myenv python=3.11 numpy scipy matplotlib ipython psutil pyopencl cupy pycuda nvidia::cuda
conda activate myenv
pip install pyopencl pycuda cupy-cuda12x pyvkfft
Troubleshooting installation
If you encounter issues, make sure you have the right combination of toolkit
and driver. Also, note that installing using pip
is cached, so if you change
your configuration (new toolkit version), you must make sure to recompile the
pycuda
and pyvkfft
packages, using e.g.:
pip install pycuda pyvkfft --no-cache
The pyvkfft-info
script can give you some information about current support
for OpenCL/CUDA, the version of the driver and the toolkit, and the detected
GPU devices.
Example
The simplest way to use pyvkfft is to use the pyvkfft.fft
interface, which will
automatically create the VkFFTApp (the FFT plans) according to the type of GPU
arrays (pycuda, pyopencl or cupy), and also cache these apps:
import pycuda.autoprimaryctx
import pycuda.gpuarray as cua
from pyvkfft.fft import fftn
import numpy as np
d0 = cua.to_gpu(np.random.uniform(0,1,(200,200)).astype(np.complex64))
# This will compute the fft to a new GPU array
d1 = fftn(d0)
# An in-place transform can also be done by specifying the destination
d0 = fftn(d0, d0)
# Or an out-of-place transform to an existing array (the destination array is always returned)
d1 = fftn(d0, d1)