Snapshot from this documentation build. SHA-256: 53ade58caa2e2c11e381bf8a94faa01df89f3ccbef4ab8ba63aa1b092f3248b4
1 module MHDKernelsCUDAExt
2
3 using MHDKernels
4 using CUDA
5
6 function __init__()
7 if CUDA.functional()
8 MHDKernels.register_backend!(:cuda, CUDABackend())
9 end
10 end
11
12 MHDKernels.device_zeros(::CUDABackend, ::Type{T}, dims::Dims) where {T} =
13 CUDA.zeros(T, dims)
14
15 # The historical native cube predates the Val-based Riemann interface and the
16 # compensated/source-aware updates. Production uses the shared KA algorithm.
17 function MHDKernels._cube_launch!(::Val{:auto}, be::CUDABackend,
18 s::MHDKernels.MHDState, dt::Real, ch::Real, decay::Real)
19 MHDKernels._cube_launch!(Val(:ka), be, s, dt, ch, decay)
20 end
21
22 function MHDKernels._cube_launch!(::Val{:raw}, ::CUDABackend,
23 ::MHDKernels.MHDState, ::Real, ::Real, ::Real)
24 error("CUDA impl=:raw is disabled: the historical native kernel is stale; " *
25 "use impl=:ka or impl=:auto for the fused KernelAbstractions kernel")
26 end
27
28 MHDKernels._radiation_fft_backend_key(x::CuArray) =
29 (CUDABackend, CUDA.device(x), CUDA.context(x))
30
31 # Transform only spatial axes. The fourth dimension is the photon/velocity
32 # batch, not a fourth spatial axis. Planning on CuArray dispatches to cuFFT;
33 # no host staging or per-channel device copies are involved.
34 function MHDKernels._radiation_forward_batch!(
35 chat::CuArray{Complex{T},4}, x::CuArray{T,4}) where {T}
36 forward, _ = MHDKernels._radiation_fft_plans(x, chat)
37 MHDKernels.FFTW.mul!(chat, forward, x)
38 chat
39 end
40
41 function MHDKernels._radiation_inverse_batch!(
42 x::CuArray{T,4}, chat::CuArray{Complex{T},4}) where {T}
43 _, inverse = MHDKernels._radiation_fft_plans(x, chat)
44 MHDKernels.FFTW.mul!(x, inverse, chat)
45 x
46 end
47
48 end # module