Multigrid inverse Laplacian (includes CombinatorialSpaces.jl v0.6.8 support)

ext/DecapodesCUDAExt.jl

@@ -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)

src/operators.jl

@@ -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
 
@@ -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)
@@ -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
 
@@ -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

src/simulation.jl

@@ -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)
 
@@ -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))
@@ -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

test/simulation.jl

@@ -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)
@@ -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.
 
@@ -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
@@ -829,4 +869,3 @@ haystack = string(gensim(LargeSum))
 @test occursin(needle, haystack)
 
 end
-