CUDA Calls

As with C backends, custom callbacks can be introduced to CUDA kernels generated by Devito. This notebook serves as a supplement to ccalls.ipynb, demonstrating how the approach showcased therein can be extended to other languages, such as CUDA.

import os
import matplotlib.pyplot as plt
import numpy as np

from scipy.signal import hilbert

from devito import TimeFunction, Eq, solve, Operator, ConditionalDimension
from devito.types.dimension import CustomDimension
from devitopro.types.enriched import NoLayers
from devitopro import *
from examples.seismic import demo_model, TimeAxis, RickerSource

#NBVAL_IGNORE_OUTPUT
so = 8
nbl = 30
model = demo_model(preset='layers-isotropic', nlayers=5,
                   shape=(301, 301), spacing=(10., 10.),
                   space_order=so, nbl=nbl)

t0, tn = -100., 1200.
dt = model.critical_dt  # Time step from model grid spacing

time_range = TimeAxis(start=t0, stop=tn, step=dt)

f0 = 0.020  # Source peak frequency is 20Hz (0.020 kHz)

src = RickerSource(name='src', grid=model.grid, f0=f0,
                   npoint=1, time_range=time_range, dtype=np.complex64)

src.data[:] = hilbert(np.real(src.data), axis=0)

# First, position source centrally in all dimensions, then set depth
src.coordinates.data[0, :] = np.array(model.domain_size) * .5
src.coordinates.data[0, -1] = 20.  # Depth is 20m

u = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=so, dtype=np.complex64)

pde = model.m * u.dt2 - u.laplace + model.damp * u.dt
stencil = Eq(u.forward, solve(pde, u.forward))

src_term = src.inject(field=u.forward, expr=src*dt**2/model.m)

#NBVAL_SKIP
# Set to the location of your Cuda installation
cufftdir = "/opt/nvhpc/math_libs"

class CuFFTCall(CCall):
    """Abstract class for C calls using CuFFT"""
    header = 'updown_cuda.h'  # Header containing CuFFT calls
    header_dirs = (os.path.abspath(""), f'{cufftdir}/include/')
    libs = ('cufft',)
    lib_dirs = (f'{cufftdir}/lib64',)


class FwdFFT(CuFFTCall):
    """Forward FFT using CuFFT"""
    name = 'fftFwd'


class InvFFT(CuFFTCall):
    """Forward FFT using CuFFT"""
    name = 'fftInv'


# Snapshotting is set up broadly as per usual
nsnaps = 30
factor = int(np.ceil(src.nt/nsnaps))

time_subsampled = ConditionalDimension('t_sub', parent=model.grid.time_dim, factor=factor)

# Upgoing and downgoing wavefield snapshots
# NoLayers disables streaming/compression for ease of plotting the wavefields in this notebook
usave_u = TimeFunction(name='usaveu', grid=model.grid, time_order=0, space_order=so,
                       save=nsnaps, time_dim=time_subsampled, dtype=np.complex64,
                       layers=NoLayers)
usave_d = TimeFunction(name='usaved', grid=model.grid, time_order=0, space_order=so,
                       save=nsnaps, time_dim=time_subsampled, dtype=np.complex64,
                       layers=NoLayers)

# Temporaries as destination for upward and downward portions of
# Fourier-transformed wavefield
tmp_u = TimeFunction(name='tmpu', grid=model.grid, space_order=so, time_order=0, dtype=np.complex64)
tmp_d = TimeFunction(name='tmpd', grid=model.grid, space_order=so, time_order=0, dtype=np.complex64)


# CustomDimensions used to apply up-down filters to the transformed wavefield
y = model.grid.dimensions[-1]
negk = CustomDimension(name='negk', symbolic_min=y.symbolic_min-so, symbolic_max=y.symbolic_max//2-1, parent=y)
posk = CustomDimension(name='posk', symbolic_min=y.symbolic_max//2+1, symbolic_max=y.symbolic_max+so, parent=y)

# Upward and downward filtering operations
upfilter = Eq(tmp_u.subs(y, posk), 0, implicit_dims=time_subsampled)
downfilter = Eq(tmp_d.subs(y, negk), 0, implicit_dims=time_subsampled)

save = [FwdFFT([u.dmap, tmp_u.dmap, tmp_d.dmap,
                u.symbolic_shape[1], model.grid.shape[-1] + 2*so,
                u.symbolic_shape[2], model.grid.time_dim], 
                implicit_dims=time_subsampled),
        upfilter, downfilter,
        InvFFT([tmp_u.dmap, tmp_d.dmap,
                u.symbolic_shape[1], model.grid.shape[-1] + 2*so,
                u.symbolic_shape[2]], 
                implicit_dims=time_subsampled),
        Eq(usave_u, tmp_u), Eq(usave_d, tmp_d)]

#NBVAL_SKIP
op_split = Operator([stencil] + save + src_term)

#NBVAL_SKIP
op_split(dt=dt)

#NBVAL_SKIP
def plot_field(field, model, nbl, title):
    clip = 0.5  # Clip value
    vmax = clip*np.amax(np.abs(field[-1]))
    vmin = -vmax

    plt.imshow(model.vp.data[nbl:-nbl, nbl:-nbl].T, cmap='Greys')
    plt.imshow(field[-1, nbl:-nbl, nbl:-nbl].T, cmap='seismic',
               vmin=vmin, vmax=vmax, alpha=0.7)
    plt.title(title)
    plt.show()


plot_field(np.real(usave_u.data) + np.real(usave_d.data), model, nbl, "Full")
plot_field(np.real(usave_u.data), model, nbl, "Upgoing")
plot_field(np.real(usave_d.data), model, nbl, "Downgoing")

Other Formats