Make streamfunctionfrompv! faster by using StaticArrays and KernelAbstractions? about geophysicalflows.jl

[navidcy]

@glwagner, S should be an Array of StaticArrays or a StaticArray of StaticArrays? I guess the former..

[navidcy]

A MWE:

using FourierFlows,  Random, FFTW, LinearAlgebra
import GeophysicalFlows.MultilayerQG
using StaticArrays, BenchmarkTools

dev=GPU()
A = ArrayType(dev)

n, L = 128, 2π
grid = TwoDGrid(dev, n, L)

nkr, nl = grid.nkr, grid.nl

nlayers = 3            # these choice of parameters give the
f0, g = 1, 1           # desired PV-streamfunction relations
H = [0.25, 0.25, 0.5]  # q1 = Δψ1 + 20ψ2 - 20ψ1,
ρ = [4.0, 5.0, 6.0]    # q2 = Δψ2 + 20ψ1 - 44ψ2 + 24ψ3,
                       # q3 = Δψ3        + 12ψ2 - 12ψ3.

T = Float64

ρ = reshape(T.(ρ), (1,  1, nlayers))
H = reshape(T.(H), (1,  1, nlayers))

g′ = T(g) * (ρ[2:nlayers] - ρ[1:nlayers-1]) ./ ρ[2:nlayers] # reduced gravity at each interface;

Fm = @. T( f0^2 / (g′ * H[2:nlayers]) )
Fp = @. T( f0^2 / (g′ * H[1:nlayers-1]) )
Fp = @. T( f0^2 / (g′ * H[1:nlayers-1]) )


function calcS_static!(S, Fp, Fm, grid)
  F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp))
  for n=1:grid.nl, m=1:grid.nkr
     k² = grid.Krsq[m, n]
    Skl = SMatrix{nlayers, nlayers}(- k²*I + F)
    S[m, n] = Skl
  end
  return nothing
end

function calcinvS_static!(invS, Fp, Fm, grid)
  F = Matrix(Tridiagonal(Fm, -([Fp; 0] + [0; Fm]), Fp))
  for n=1:grid.nl, m=1:grid.nkr
    k² = grid.Krsq[m, n]
    if k² == 0
      k² = 1
    end
    Skl = - k²*I + F
    invS[m, n] = SMatrix{nlayers, nlayers}(I / Skl)
  end
  invS[1, 1] = SMatrix{nlayers, nlayers}(zeros(dev, T, (nlayers, nlayers)))
  return nothing
end

## Construct S and invS the as 4D arrays
S = Array{T}(undef, (nkr, nl, nlayers, nlayers))
MultilayerQG.calcS!(S, Fp, Fm, grid)

invS = Array{T}(undef, (nkr, nl, nlayers, nlayers))
MultilayerQG.calcinvS!(invS, Fp, Fm, grid)

S, invS = A(S), A(invS)     # convert S and invS to appropriate ArrayType

## Construct S and invS the as 2D arrays of 2D StaticArrays
Skl_static = SMatrix{nlayers, nlayers}(zeros(dev, T, (nlayers, nlayers)))

S_static = Array{typeof(Skl_static), 2}(undef, (nkr, nl))
calcS_static!(S_static, Fp, Fm, grid)

invS_static = Array{typeof(Skl_static), 2}(undef, (nkr, nl))
calcinvS_static!(invS_static, Fp, Fm, grid)

S_static, invS_static = A(S_static), A(invS_static)     # convert S and invS to appropriate ArrayType

## Benchmark the two methods
function constructtestfields_3layer(gr)
  x, y = gridpoints(gr)
  k₀, l₀ = 2π/gr.Lx, 2π/gr.Ly # fundamental wavenumbers

  # a set of streafunctions ψ1, ψ2, ψ3, ...
  ψ1 = @. 1e-3 * ( 1/4*cos(2k₀*x)*cos(5l₀*y) + 1/3*cos(3k₀*x)*cos(3l₀*y) )
  ψ2 = @. 1e-3 * (     cos(3k₀*x)*cos(4l₀*y) + 1/2*cos(4k₀*x)*cos(2l₀*y) )
  ψ3 = @. 1e-3 * (     cos(1k₀*x)*cos(3l₀*y) + 1/2*cos(2k₀*x)*cos(2l₀*y) )
  
  Δψ1 = @. -1e-3 * ( 1/4*((2k₀)^2+(5l₀)^2)*cos(2k₀*x)*cos(5l₀*y) + 1/3*((3k₀)^2+(3l₀)^2)*cos(3k₀*x)*cos(3l₀*y) )
  Δψ2 = @. -1e-3 * (     ((3k₀)^2+(4l₀)^2)*cos(3k₀*x)*cos(4l₀*y) + 1/2*((4k₀)^2+(2l₀)^2)*cos(4k₀*x)*cos(2l₀*y) )
  Δψ3 = @. -1e-3 * (     ((1k₀)^2+(3l₀)^2)*cos(1k₀*x)*cos(3l₀*y) + 1/2*((2k₀)^2+(2l₀)^2)*cos(2k₀*x)*cos(2l₀*y) )

  # ... their corresponding PVs q1, q2, q3,
  q1 = @. Δψ1 + 20ψ2 - 20ψ1
  q2 = @. Δψ2 + 20ψ1 - 44ψ2 + 24ψ3
  q3 = @. Δψ3        + 12ψ2 - 12ψ3

  return ψ1, ψ2, ψ3, q1, q2, q3
end

ψ1, ψ2, ψ3, q1, q2, q3 = constructtestfields_3layer(grid)

qh = zeros(dev, Complex{T}, (nkr, nl, nlayers))
ψh = zeros(dev, Complex{T}, (nkr, nl, nlayers))
ψh_static = zeros(dev, Complex{T}, (nkr, nl, nlayers))

qh[:, :, 1] .= rfft(q1)
qh[:, :, 2] .= rfft(q2)
qh[:, :, 3] .= rfft(q3)

# repeat here to get invS as arg and not params
function streamfunctionfrompv!(ψh, qh, invS, grid)
  for j=1:grid.nl, i=1:grid.nkr
    @views ψh[i, j, :] .= invS[i, j, :, :] * qh[i, j, :]
  end
end

function streamfunctionfrompv_static!(ψh, qh, invS_static, grid)
  for j=1:grid.nl, i=1:grid.nkr
    @views ψh[i, j, :] .= invS_static[i, j] * qh[i, j, :]
  end
end

@btime streamfunctionfrompv!(ψh, qh, invS, grid)

@btime streamfunctionfrompv_static!(ψh_static, qh, invS_static, grid)

isapprox(ψh, ψh_static)

outputs:

julia> @btime streamfunctionfrompv!(ψh, qh, invS, grid)
  518.804 μs (33280 allocations: 2.54 MiB)

julia> @btime streamfunctionfrompv_static!(ψh_static, qh, invS_static, grid)
  42.102 μs (0 allocations: 0 bytes)

julia> isapprox(ψh, ψh_static)
true

[navidcy]

with dev=GPU() btw I get:

julia> @btime streamfunctionfrompv!(ψh, qh, invS, grid)
  3.458 s (1886180 allocations: 77.53 MiB)

julia> @btime streamfunctionfrompv_static!(ψh_static, qh, invS_static, grid)
  891.316 ms (357760 allocations: 15.23 MiB)

[navidcy]

But still no KernelAbstractions....