Skip to content

Commit

Permalink
Merge pull request #276 from AlgebraicJulia/gr/inv_lap
Browse files Browse the repository at this point in the history
Multigrid inverse Laplacian (includes CombinatorialSpaces.jl v0.6.8 support)
  • Loading branch information
jpfairbanks authored Dec 4, 2024
2 parents dc75301 + 4ebafab commit 8eec5c9
Show file tree
Hide file tree
Showing 4 changed files with 89 additions and 28 deletions.
18 changes: 9 additions & 9 deletions ext/DecapodesCUDAExt.jl
Original file line number Diff line number Diff line change
Expand Up @@ -95,21 +95,21 @@ function dec_cu_mat_dual_differential(k::Int, sd::HasDeltaSet)
end

function dec_cu_pair_wedge_product(::Type{Tuple{k,0}}, sd::HasDeltaSet) where {k}
val_pack = dec_p_wedge_product(Tuple{0,k}, sd, Val{:CUDA})
((y, α, g) -> dec_c_wedge_product!(Tuple{0,k}, y, g, α, val_pack, Val{:CUDA}),
(α, g) -> dec_c_wedge_product(Tuple{0,k}, g, α, val_pack, Val{:CUDA}))
val_pack = cache_wedge(Tuple{0,k}, sd, Val{:CUDA})
((y, α, g) -> dec_c_wedge_product!(Tuple{0,k}, y, g, α, val_pack[1], val_pack[2]),
(α, g) -> dec_c_wedge_product(Tuple{0,k}, g, α, val_pack))
end

function dec_cu_pair_wedge_product(::Type{Tuple{0,k}}, sd::HasDeltaSet) where {k}
val_pack = dec_p_wedge_product(Tuple{0,k}, sd, Val{:CUDA})
((y, f, β) -> dec_c_wedge_product!(Tuple{0,k}, y, f, β, val_pack, Val{:CUDA}),
(f, β) -> dec_c_wedge_product(Tuple{0,k}, f, β, val_pack, Val{:CUDA}))
val_pack = cache_wedge(Tuple{0,k}, sd, Val{:CUDA})
((y, f, β) -> dec_c_wedge_product!(Tuple{0,k}, y, f, β, val_pack[1], val_pack[2]),
(f, β) -> dec_c_wedge_product(Tuple{0,k}, f, β, val_pack))
end

function dec_cu_pair_wedge_product(::Type{Tuple{1,1}}, sd::HasDeltaSet2D)
val_pack = dec_p_wedge_product(Tuple{1,1}, sd, Val{:CUDA})
((y, α, β) -> dec_c_wedge_product!(Tuple{1,1}, y, α, β, val_pack, Val{:CUDA}),
(α, β) -> dec_c_wedge_product(Tuple{1,1}, α, β, val_pack, Val{:CUDA}))
val_pack = cache_wedge(Tuple{1,1}, sd, Val{:CUDA})
((y, α, β) -> dec_c_wedge_product!(Tuple{1,1}, y, α, β, val_pack[1], val_pack[2]),
(α, β) -> dec_c_wedge_product(Tuple{1,1}, α, β, val_pack))
end

function dec_pair_wedge_product(::Type{Tuple{0,0}}, sd::HasDeltaSet)
Expand Down
32 changes: 24 additions & 8 deletions src/operators.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,14 @@ using SparseArrays

function default_dec_cu_matrix_generate() end;

function default_dec_matrix_generate(fs::PrimalGeometricMapSeries, my_symbol::Symbol, hodge::DiscreteHodge)
op = @match my_symbol begin
# Inverse Laplacians
:Δ₀⁻¹ => dec_Δ⁻¹(Val{0}, fs)
_ => default_dec_matrix_generate(finest_mesh(fs), my_symbol, hodge)
end
end

function default_dec_matrix_generate(sd::HasDeltaSet, my_symbol::Symbol, hodge::DiscreteHodge)
op = @match my_symbol begin

Expand Down Expand Up @@ -54,8 +62,11 @@ function default_dec_matrix_generate(sd::HasDeltaSet, my_symbol::Symbol, hodge::
:ℒ₁ => ℒ_dd(Tuple{1,1}, sd)

# Dual Laplacians
:Δᵈ₀ => Δᵈ(Val{0},sd)
:Δᵈ₁ => Δᵈ(Val{1},sd)
:Δᵈ₀ => Δᵈ(Val{0}, sd)
:Δᵈ₁ => Δᵈ(Val{1}, sd)

# Inverse Laplacians
:Δ₀⁻¹ => dec_inv_lap_solver(Val{0}, sd)

# Musical Isomorphisms
:♯ => dec_♯_p(sd)
Expand Down Expand Up @@ -105,20 +116,20 @@ function dec_mat_dual_differential(k::Int, sd::HasDeltaSet)
end

function dec_pair_wedge_product(::Type{Tuple{k,0}}, sd::HasDeltaSet) where {k}
val_pack = dec_p_wedge_product(Tuple{0,k}, sd)
((y, α, g) -> dec_c_wedge_product!(Tuple{0,k}, y, g, α, val_pack),
val_pack = cache_wedge(Tuple{0,k}, sd, Val{:CPU})
((y, α, g) -> dec_c_wedge_product!(Tuple{0,k}, y, g, α, val_pack[1], val_pack[2]),
(α, g) -> dec_c_wedge_product(Tuple{0,k}, g, α, val_pack))
end

function dec_pair_wedge_product(::Type{Tuple{0,k}}, sd::HasDeltaSet) where {k}
val_pack = dec_p_wedge_product(Tuple{0,k}, sd)
((y, f, β) -> dec_c_wedge_product!(Tuple{0,k}, y, f, β, val_pack),
val_pack = cache_wedge(Tuple{0,k}, sd, Val{:CPU})
((y, f, β) -> dec_c_wedge_product!(Tuple{0,k}, y, f, β, val_pack[1], val_pack[2]),
(f, β) -> dec_c_wedge_product(Tuple{0,k}, f, β, val_pack))
end

function dec_pair_wedge_product(::Type{Tuple{1,1}}, sd::HasDeltaSet2D)
val_pack = dec_p_wedge_product(Tuple{1,1}, sd)
((y, α, β) -> dec_c_wedge_product!(Tuple{1,1}, y, α, β, val_pack),
val_pack = cache_wedge(Tuple{1,1}, sd, Val{:CPU})
((y, α, β) -> dec_c_wedge_product!(Tuple{1,1}, y, α, β, val_pack[1], val_pack[2]),
(α, β) -> dec_c_wedge_product(Tuple{1,1}, α, β, val_pack))
end

Expand Down Expand Up @@ -146,6 +157,11 @@ function dec_avg₀₁(sd::HasDeltaSet)
(avg_mat, x -> avg_mat * x)
end

function dec_inv_lap_solver(::Type{Val{0}}, sd::HasDeltaSet)
inv_lap = LinearAlgebra.factorize(∇²(0, sd))
x -> inv_lap \ x
end

function default_dec_generate(sd::HasDeltaSet, my_symbol::Symbol, hodge::DiscreteHodge=GeometricHodge())

op = @match my_symbol begin
Expand Down
16 changes: 11 additions & 5 deletions src/simulation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -681,12 +681,12 @@ Base.showerror(io::IO, e::UnsupportedStateeltypeException) = print(io, "Decapode
"""
gensim(user_d::SummationDecapode, input_vars::Vector{Symbol}; dimension::Int=2, stateeltype::DataType = Float64, code_target::AbstractGenerationTarget = CPUTarget(), preallocate::Bool = true)
Generates the entire code body for the simulation function. The returned simulation function can then be combined with a mesh, provided by `CombinatorialSpaces`, and a function describing symbol
Generates the entire code body for the simulation function. The returned simulation function can then be combined with a mesh, provided by `CombinatorialSpaces`, and a function describing symbol
to operator mappings to return a simulator that can be used to solve the represented equations given initial conditions.
**Arguments:**
`user_d`: The user passed Decapode for which simulation code will be generated. (This is not modified)
`user_d`: The user passed Decapode for which simulation code will be generated. (This is not modified)
`input_vars` is the collection of variables whose values are known at the beginning of the simulation. (Defaults to all state variables and literals in the Decapode)
Expand All @@ -699,8 +699,10 @@ to operator mappings to return a simulator that can be used to solve the represe
`code_target`: The intended architecture target for the generated code. (Defaults to `CPUTarget()`)(Use `CUDATarget()` for NVIDIA CUDA GPUs)
`preallocate`: Enables(`true`)/disables(`false`) pre-allocated caches for intermediate computations. Some functions, such as those that determine Jacobian sparsity patterns, or perform auto-differentiation, may require this to be disabled. (Defaults to `true`)
`multigrid`: Enables multigrid methods during code generation. If `true`, then the function produced by `gensim` will expect a `PrimalGeometricMapSeries`. (Defaults to `false`)
"""
function gensim(user_d::SummationDecapode, input_vars::Vector{Symbol}; dimension::Int=2, stateeltype::DataType = Float64, code_target::AbstractGenerationTarget = CPUTarget(), preallocate::Bool = true)
function gensim(user_d::SummationDecapode, input_vars::Vector{Symbol}; dimension::Int=2, stateeltype::DataType = Float64, code_target::AbstractGenerationTarget = CPUTarget(), preallocate::Bool = true, multigrid::Bool = false)

(dimension == 1 || dimension == 2) ||
throw(UnsupportedDimensionException(dimension))
Expand Down Expand Up @@ -754,10 +756,14 @@ function gensim(user_d::SummationDecapode, input_vars::Vector{Symbol}; dimension
func_defs = compile_env(gen_d, dec_matrices, contracted_dec_operators, code_target)
vect_defs = compile_var(alloc_vectors)

prologue = quote end
multigrid && push!(prologue.args, :(mesh = finest_mesh(mesh)))

quote
(mesh, operators, hodge=GeometricHodge()) -> begin
$func_defs
$cont_defs
$prologue
$vect_defs
f(__du__, __u__, __p__, __t__) = begin
$vars
Expand Down
51 changes: 45 additions & 6 deletions test/simulation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,11 @@ using LinearAlgebra
using MLStyle
using OrdinaryDiffEq
using Test
using Random
Point3D = Point3{Float64}

import Decapodes: default_dec_matrix_generate

flatten(vfield::Function, mesh) = (mesh, DualVectorField(vfield.(mesh[triangle_center(mesh),:dual_point])))

function test_hodge(k, sd::HasDeltaSet, hodge)
Expand Down Expand Up @@ -757,6 +761,46 @@ end

end

@testset "Multigrid" begin
s = triangulated_grid(1,1,1/4,sqrt(3)/2*1/4,Point3D)

series = PrimalGeometricMapSeries(s, binary_subdivision_map, 4);

our_mesh = finest_mesh(series)
lap = ∇²(0,our_mesh);

Random.seed!(1337)
b = lap*rand(nv(our_mesh));

inv_lap = @decapode begin
U::Form0
∂ₜ(U) == Δ₀⁻¹(U)
end

function generate(fs, my_symbol; hodge=DiagonalHodge())
op = @match my_symbol begin
_ => default_dec_matrix_generate(fs, my_symbol, hodge)
end
end

sim = eval(gensim(inv_lap))
sim_mg = eval(gensim(inv_lap; multigrid=true))

f = sim(our_mesh, generate);
f_mg = sim_mg(series, generate);

u = ComponentArray(U=b)
du = similar(u)

# Regular mesh
f(du, u, 0, ())
@test norm(lap*du.U-b)/norm(b) < 1e-15

# Multigrid
f_mg(du, u, 0, ())
@test norm(lap*du.U-b)/norm(b) < 1e-6
end

@testset "Allocations" begin
# Test the heat equation Decapode has expected memory allocation.

Expand All @@ -781,11 +825,7 @@ for prealloc in [false, true]
nallocs = @allocations f(du, u₀, p, (0,1.0))
bytes = @allocated f(du, u₀, p, (0,1.0))

if VERSION < v"1.11"
@test (nallocs, bytes) == (prealloc ? (3, 80) : (5, 400))
else
@test (nallocs, bytes) == (prealloc ? (2, 32) : (6, 352))
end
@test (nallocs, bytes) <= (prealloc ? (6, 80) : (6, 400))
end

end
Expand Down Expand Up @@ -829,4 +869,3 @@ haystack = string(gensim(LargeSum))
@test occursin(needle, haystack)

end

0 comments on commit 8eec5c9

Please sign in to comment.