Rework: Extensions

[albert-de-montserrat]

Major rework

This PR involves a decent full rework of the internal design of JustRelax.jl.

Main changes

1. Two-step initialization has been deprecated in favour of using specialized two- and three-dimensional submodules. These submodules include essentially the same code, but they are compiled for the two and three dimensional backends of ParallelStencil. For example, the old initialization

using JustRelax # first stage
model = PS_Setup(:cpu, Float64, 2) 
environment!(model) # second compilation stage

is deprecated in favour of

using JustRelax, JustRelax.JustRelax2D

or

using JustRelax, JustRelax.JustRelax3D

2. Extensions are used to switch from the CPU to supported GPU backends. This simplifies the setup of the package by removing lots of hacky and inefficient code, and improves pre-compilation. The user now just needs to choose between CPUBackend, CUDABackend,AMDGPUBackend to instantiate the basic JustRelax objects.

3. To ease a bit the usage of JustRelax and improve readibility, all the functions that required the @parallel macro are not exported anymore. Instead, wrappers around them are exported. E.g.

compute_ρg!(ρg[2], phase_ratios, rheology, args)

substitutes the old

@parallel (@idx ni) compute_ρg!(ρg[2], phase_ratios.center, rheology, args)

[albert-de-montserrat]

2D shear heating benchmark, 64x64, 2 threads, 3 time steps. This branch:

t1 = 13.703740 seconds # (20.70 M allocations: 1.396 GiB, 4.58% gc time, 74.67% compilation time); warm up run (includes compilation)
t2 = 4.833359 seconds # (4.83 M allocations: 412.365 MiB, 0.33% gc time); second run (no compilation)

#main branch

t1 = 15.079457 seconds # (24.12 M allocations: 1.645 GiB, 6.54% gc time, 72.40% compilation time); warm up run (includes compilation) 
t2 = 5.245096 seconds # (5.15 M allocations: 449.310 MiB, 0.31% gc time); second run (no compilation)