CuPy and PyCUDA comparison
Note that mixing pycuda and cupy isn’t a very good idea, as the handling of CUDA contexts is different… But this works as far as demonstrating CuPy and PyCUDA give the same results.
[1]:
%matplotlib notebook
import numpy as np
try:
from scipy.datasets import ascent
except ImportError:
from scipy.misc import ascent
from scipy.fft import fftn, ifftn, rfftn, irfftn, fftshift
from matplotlib.colors import LogNorm
import matplotlib.pyplot as plt
import cupy as cp
# This forces the CuPy context init
cupy_array = cp.zeros((4,4))
import pycuda
import pycuda.autoinit
import pycuda.gpuarray as cua
from pyvkfft.cuda import *
[2]:
# Inplace, CuPy only
for dtype in [np.complex64, np.float32]:
a = ascent().astype(dtype)
app = VkFFTApp(a.shape, a.dtype, r2c=(a.dtype==np.float32))
print(dtype, app.r2c)
a_cupy = cp.asarray(a)
plt.figure(figsize=(9.,3))
plt.subplot(131)
plt.imshow(abs(a))
plt.colorbar()
plt.subplot(132)
if a.dtype == np.float32:
b = rfftn(a[...,:-2])
ax = [0]
else:
b = fftn(a)
ax = None
plt.imshow(abs(fftshift(b, axes=ax)), norm=LogNorm())
plt.colorbar()
plt.title("scipy")
plt.subplot(133)
a_cupy = app.fft(a_cupy)
plt.imshow(abs(fftshift(a_cupy.get(), axes=ax)), norm=LogNorm())
plt.colorbar()
plt.title("cupy")
plt.tight_layout()
<class 'numpy.complex64'> False
<class 'numpy.float32'> True
[5]:
# Inplace
for dtype in [np.complex64, np.float32]:
a = ascent().astype(dtype)
app = VkFFTApp(a.shape, a.dtype, r2c=(a.dtype==np.float32))
print(dtype, app.r2c)
a_pycuda = cua.to_gpu(a)
a_cupy = cp.asarray(a)
plt.figure(figsize=(9.,3))
plt.subplot(141)
plt.imshow(abs(a))
plt.colorbar()
plt.subplot(142)
if a.dtype == np.float32:
ax = [0]
plt.imshow(abs(fftshift(rfftn(a), axes=ax)), norm=LogNorm())
else:
ax=None
plt.imshow(abs(fftshift(fftn(a))), norm=LogNorm())
plt.colorbar()
plt.title("scipy")
plt.subplot(143)
plt.imshow(abs(fftshift(app.fft(a_pycuda).get(), axes=ax)), norm=LogNorm())
plt.colorbar()
plt.title("pycuda")
plt.subplot(144)
plt.imshow(abs(fftshift(app.fft(a_cupy).get(), axes=ax)), norm=LogNorm())
plt.colorbar()
plt.title("cupy")
plt.tight_layout()
print(np.allclose(a_pycuda.get(), a_cupy.get()))
app.ifft(a_pycuda)
app.ifft(a_cupy)
print(np.allclose(a_pycuda.get(), a_cupy.get()))
<class 'numpy.complex64'> False
True
True
<class 'numpy.float32'> True
True
True
[4]:
# Out-of-place
for dtype in [np.complex64, np.float32]:
print(dtype)
a = ascent().astype(dtype)
app = VkFFTApp(a.shape, a.dtype, r2c=(a.dtype==np.float32), inplace=False)
a_pycuda = cua.to_gpu(a)
a_cupy = cp.asarray(a)
if dtype == np.complex64:
b_pycuda = cua.empty_like(a_pycuda)
b_cupy = cp.empty_like(a_cupy)
else:
ny, nx = a.shape
nx2 = nx //2 + 1
b_pycuda = cua.empty((ny,nx2), np.complex64)
b_cupy = cp.empty((ny,nx2), np.complex64)
plt.figure(figsize=(9.,3))
plt.subplot(141)
plt.imshow(abs(a))
plt.colorbar()
plt.subplot(142)
if a.dtype == np.float32:
ax = [0]
plt.imshow(abs(fftshift(rfftn(a), axes=ax)), norm=LogNorm())
else:
ax = None
plt.imshow(abs(fftshift(fftn(a))), norm=LogNorm())
plt.colorbar()
plt.title("scipy")
plt.subplot(143)
plt.imshow(abs(fftshift(app.fft(a_pycuda, b_pycuda).get(), axes=ax)), norm=LogNorm())
plt.colorbar()
plt.title("pycuda")
plt.subplot(144)
plt.imshow(abs(fftshift(app.fft(a_cupy, b_cupy).get(), axes=ax)), norm=LogNorm())
plt.colorbar()
plt.title("cupy")
plt.tight_layout()
print(np.allclose(b_pycuda.get(), b_cupy.get()))
app.ifft(b_pycuda, a_pycuda)
app.ifft(b_cupy, a_cupy)
print(np.allclose(a_pycuda.get(), a_cupy.get()))
<class 'numpy.complex64'>
True
True
<class 'numpy.float32'>
True
True
[ ]: