WaterLily + AutoDiff workflow

[weymouth]

The goal of this PR is to enable auto-diff with WaterLily, ala issue #52. This PR is IN-PROGRESS!

A minor issue is that we have certain parts of the code which declare types which should not. I've removed some of these in the first commit.

The second and far more serious issue is that of the field arrays within Flow and Poisson. These are pre-allocated, and this is incompatible with ForwardDiff, which needs to be able to allocate Dual arrays for variables "downstream" of the derivative being taken. I've included an example "test.jl" file which fails with ForwardDiff when the simulation is pre-allocated as normal OR when pre-allocated with a Dual type.

The last issue is KernelAbstractions. It's possible that this will not be an issue once the types are sorted out, but the error messages are so bad that it's impossible to debug. I've changed @loop to just use @SIMD for the time being.

[b-fg]

I can reproduce the error too. I am looking into it.

[b-fg]

Nice, works as expected!

[b-fg]

KA with CPU backend is working. However, trying to use KA with CuArray is failing even if using: θ = Ref(π/36) |> cudaconvert, which returns a CuRefValue. Without AD, if fails when trying to modify the CuRefValue, ie θ[]=0.5:

LoadError: MethodError: no method matching setindex!(::CUDA.CuRefValue{Float64}, ::Float64)

Closest candidates are:
  setindex!(::Ref, ::Any, ::CartesianIndex{0})
   @ Base multidimensional.jl:1933

I guess it means that we cannot directly modify the pointer value in GPU memory? Maybe @maleadt has an idea on how to circumvent this? KA on GPU with θ = Ref(π/36) |> cudaconvert (no AD) works well though.

On the other hand, when using AD and without modifying the CuRefValue

T = typeof(Tag(step_force!, Float64)) # make a tag
θAD = Ref(Dual{T}(π/36, 0.)) |> cudaconvert # wrap the Dual parameter in a CuRefValue

there seems to be a problem when executing the measure kernel of Body, on BC! https://github.com/weymouth/WaterLily.jl/blob/d9161b7dcc262e4744ac88649b8d21b7310f354d/src/Body.jl#L46
Maybe related to the BC! problem on GPU you mentioned before @weymouth ?

[weymouth]

That Ref problem on GPUs doesn't have anything to do with WaterLily in particular. We can play around with really small MWEs and see what works. Maybe post something on Slack.

The second problem sounds like something we need to fix in the code base. Does that bit work on the CPU as is?

[b-fg]

Interestingly, with the new generalization of @loop to use @simd or KA based on Threads.nthreads(), I get a speedup when using a single thread compared to the bare @simd implementation, even though there are many more allocations.
A)

macro loop(args...)
    ex,_,itr = args
    _,I,R = itr.args
    return quote
        @simd for $I ∈ $R
            @fastmath @inbounds $ex
        end
    end |> esc
end

B)

macro loop(args...)
    ex,_,itr = args
    _,I,R = itr.args; sym = []
    grab!(sym,ex)     # get arguments and replace composites in `ex`
    setdiff!(sym,[I]) # don't want to pass I as an argument
    @gensym(kern, kern_) # generate unique kernel function names for serial and KA execution
    return quote
        function $kern($(rep.(sym)...),::Val{1})
            @simd for $I ∈ $R
                @fastmath @inbounds $ex
            end
        end
        @kernel function $kern_($(rep.(sym)...),@Const(I0)) # replace composite arguments
            $I = @index(Global,Cartesian)
            $I += I0
            @fastmath @inbounds $ex
        end
        function $kern($(rep.(sym)...),_)
            $kern_(get_backend($(sym[1])),64)($(sym...),$R[1]-oneunit($R[1]),ndrange=size($R))
        end
        $kern($(sym...),Val{Threads.nthreads()}()) # dispatch to SIMD for -t 1, or KA otherwise
    end |> esc
end

resulting in the following benchmarks of @btime sim_step!($sim)

A) 776.913 μs (180 allocations: 250.23 KiB)
B) 656.911 μs (1376 allocations: 269.77 KiB)

[b-fg]

The simple SIMD version of @loop

macro loop(args...)
    ex,_,itr = args
    _,I,R = itr.args
    return quote
        @simd for $I ∈ $R
            @fastmath @inbounds $ex
        end
    end |> esc
end

Still shows 62 allocations on sim_step! for me. How did you write it to get 0 allocations @weymouth? Or was that at mom_step!?

EDIT: Nvm, it is 0 in mom_step! indeed.

[b-fg]

Interestingly enough, the simple @loop version produces 0 allocations, but is slower.
Simple @loop

576.174 μs (0 allocations: 0 bytes)

Current "functional" @loop

434.004 μs (486 allocations: 7.84 KiB)

Since the functional @loop is faster (?) and does not incur many allocations (these are coming from grab! and rep), I think the current form is acceptable.

[weymouth]

@b-fg , What is the test case you are using for timing in those comments above?

[b-fg]

This one

using WaterLily, StaticArrays, BenchmarkTools

function make_sim(θ;L=32,U=1,Re=100,mem=Array)
    function map(x,t)
        s,c = sincos(θ)
        SA[c -s; s c]*(x-SA[L,L])
    end
    function sdf(ξ,t)
        p = ξ-SA[0,clamp(ξ[1],-L/2,L/2)]
        √(p'*p)-2
    end
    Simulation((2L,2L),(U,0),L,ν=U*L/Re,body=AutoBody(sdf,map),T=Float32,mem=mem)
end

sim = make_sim(π/36; mem=Array)
sim_step!(sim)
@btime mom_step!($sim.flow, $sim.pois)

[weymouth]

I can replicate that on my Windows PC at home. With sim = make_sim(0f0) I'm getting

796.000 μs (0 allocations: 0 bytes)

vs

626.200 μs (413 allocations: 6.70 KiB)

I also checked the impact of array size using sim = make_sim(0f0,L=256) which gives

31.799 ms (0 allocations: 0 bytes)
vs

27.886 ms (264 allocations: 4.38 KiB)

So the amount of allocations does not grow with array size and is consistently a little faster. Certainly not an issue.

[b-fg]

Are tests passing locally for CuArrays too? What else is needed for this PR?

[codecov]

Codecov Report

Attention: Patch coverage is 95.91837% with 2 lines in your changes missing coverage. Please review.

Project coverage is 94.21%. Comparing base (f6c5b34) to head (7750fb6).
Report is 2 commits behind head on master.

❗ Current head 7750fb6 differs from pull request most recent head 9bac6a0

Please upload reports for the commit 9bac6a0 to get more accurate results.

Files	Patch %	Lines
src/Poisson.jl	83.33%	1 Missing ⚠️
src/util.jl	93.75%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #129      +/-   ##
==========================================
- Coverage   96.02%   94.21%   -1.82%     
==========================================
  Files          12       12              
  Lines         503      501       -2     
==========================================
- Hits          483      472      -11     
- Misses         20       29       +9     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[weymouth]

How can the code coverage be worse? I didn't change the number of lines or functions? I reduced the number of arguments in two functions!!?

[b-fg]

It's also complaining about src/util.jl#L105, which you did not touch at all... Nvm, I think we are good.

[b-fg]

Doing some review, shouldn't the T(0) be zero(T) in
https://github.com/weymouth/WaterLily.jl/blob/b97d700dea197c2683f70d68753edaa591329bbf/src/WaterLily.jl#L71

[b-fg]

Went over it once again, all looks good except for the 2 minor details above!

[vchuravy]

@michel2323 has been working on adding Enzyme support to KernelAbstractions which together with Checkpointing.jl adds reverse mode capabilities.

We just gave a tutorial on using Enzyme for codes like this: https://dj4earth.github.io/MPE24 and in particular
https://dj4earth.github.io/MPE24/Burgers_tutorial/burgers_tutorial.html

We are still working through some minutia for full GPU support, but we are getting close.

[b-fg]

Thanks @vchuravy! That's very useful indeed. Currently ForwardDiff works as expected on CPUs, but yields a fairly cryptic error message on GPUs. Eventually I think that we will move to Enzyme, probably when the GPU support under KA is ready. And thanks for working on that! It will truly bring this solver to a whole new feature level :)