From 1a6189791d7b528d3ceb3d01b59a99fb86af9ab2 Mon Sep 17 00:00:00 2001 From: JordiManyer Date: Mon, 11 Nov 2024 22:54:53 +1100 Subject: [PATCH] Added distributed autodiff --- Project.toml | 2 + src/Autodiff.jl | 105 +++++++++++++++++++++++++++++++ src/GridapDistributed.jl | 3 + test/AutodiffTests.jl | 45 +++++++++++++ test/PLaplacianTests.jl | 18 +++--- test/TestApp/src/TestApp.jl | 1 + test/mpi/runtests_np4_body.jl | 3 + test/sequential/AutodiffTests.jl | 7 +++ 8 files changed, 176 insertions(+), 8 deletions(-) create mode 100644 src/Autodiff.jl create mode 100644 test/AutodiffTests.jl create mode 100644 test/sequential/AutodiffTests.jl diff --git a/Project.toml b/Project.toml index 87919ed..7221a7d 100644 --- a/Project.toml +++ b/Project.toml @@ -6,6 +6,7 @@ version = "0.4.5" [deps] BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Gridap = "56d4f2e9-7ea1-5844-9cf6-b9c51ca7ce8e" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195" @@ -17,6 +18,7 @@ WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192" [compat] BlockArrays = "1" FillArrays = "1" +ForwardDiff = "0.10" Gridap = "0.18.7" LinearAlgebra = "1" MPI = "0.16, 0.17, 0.18, 0.19, 0.20" diff --git a/src/Autodiff.jl b/src/Autodiff.jl new file mode 100644 index 0000000..eedbee2 --- /dev/null +++ b/src/Autodiff.jl @@ -0,0 +1,105 @@ + +function Fields.gradient(f::Function,uh::DistributedCellField) + fuh = f(uh) + FESpaces._gradient(f,uh,fuh) +end + +function FESpaces._gradient(f,uh,fuh::DistributedDomainContribution) + local_terms = map(r -> DomainContribution(), get_parts(fuh)) + local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh))) + for local_trians in local_domains + g = FESpaces._change_argument(gradient,f,local_trians,uh) + cell_u = map(get_cell_dof_values,local_views(uh)) + cell_id = map(FESpaces._compute_cell_ids,local_views(uh),local_trians) + cell_grad = distributed_autodiff_array_gradient(g,cell_u,cell_id) + map(add_contribution!,local_terms,local_trians,cell_grad) + end + DistributedDomainContribution(local_terms) +end + +function Fields.jacobian(f::Function,uh::DistributedCellField) + fuh = f(uh) + FESpaces._jacobian(f,uh,fuh) +end + +function FESpaces._jacobian(f,uh,fuh::DistributedDomainContribution) + local_terms = map(r -> DomainContribution(), get_parts(fuh)) + local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh))) + for local_trians in local_domains + g = FESpaces._change_argument(jacobian,f,local_trians,uh) + cell_u = map(get_cell_dof_values,local_views(uh)) + cell_id = map(FESpaces._compute_cell_ids,local_views(uh),local_trians) + cell_grad = distributed_autodiff_array_jacobian(g,cell_u,cell_id) + map(add_contribution!,local_terms,local_trians,cell_grad) + end + DistributedDomainContribution(local_terms) +end + +function FESpaces._change_argument(op,f,trian,uh::DistributedCellField) + spaces = map(get_fe_space,local_views(uh)) + function g(cell_u) + cf = DistributedCellField( + map(CellField,spaces,cell_u), + get_triangulation(uh), + ) + cell_grad = f(cf) + map(get_contribution,local_views(cell_grad),trian) + end + g +end + +# autodiff_array_xxx + +function distributed_autodiff_array_gradient(a,i_to_x) + dummy_forwarddiff_tag = ()->() + i_to_xdual = map(i_to_x) do i_to_x + lazy_map(DualizeMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x) + end + i_to_ydual = a(i_to_xdual) + i_to_result = map(i_to_ydual,i_to_x) do i_to_ydual,i_to_x + i_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x) + lazy_map(AutoDiffMap(ForwardDiff.gradient),i_to_ydual,i_to_x,i_to_cfg) + end + return i_to_result +end + +function distributed_autodiff_array_jacobian(a,i_to_x) + dummy_forwarddiff_tag = ()->() + i_to_xdual = map(i_to_x) do i_to_x + lazy_map(DualizeMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x) + end + i_to_ydual = a(i_to_xdual) + i_to_result = map(i_to_ydual,i_to_x) do i_to_ydual,i_to_x + i_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x) + lazy_map(AutoDiffMap(ForwardDiff.jacobian),i_to_ydual,i_to_x,i_to_cfg) + end + i_to_result +end + +function distributed_autodiff_array_gradient(a,i_to_x,j_to_i) + dummy_forwarddiff_tag = ()->() + i_to_xdual = map(i_to_x) do i_to_x + lazy_map(DualizeMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x) + end + j_to_ydual = a(i_to_xdual) + j_to_result = map(i_to_x,j_to_i,j_to_ydual) do i_to_x,j_to_i,j_to_ydual + j_to_x = lazy_map(Reindex(i_to_x),j_to_i) + j_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),j_to_x) + lazy_map(AutoDiffMap(ForwardDiff.gradient),j_to_ydual,j_to_x,j_to_cfg) + end + return j_to_result +end + +function distributed_autodiff_array_jacobian(a,i_to_x,j_to_i) + dummy_forwarddiff_tag = ()->() + i_to_xdual = map(i_to_x) do i_to_x + lazy_map(DualizeMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x) + end + j_to_ydual = a(i_to_xdual) + j_to_result = map(i_to_x,j_to_i,j_to_ydual) do i_to_x,j_to_i,j_to_ydual + j_to_x = lazy_map(Reindex(i_to_x),j_to_i) + j_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),j_to_x) + lazy_map(AutoDiffMap(ForwardDiff.jacobian),j_to_ydual,j_to_x,j_to_cfg) + end + j_to_result +end diff --git a/src/GridapDistributed.jl b/src/GridapDistributed.jl index 8c20a97..3e1832f 100644 --- a/src/GridapDistributed.jl +++ b/src/GridapDistributed.jl @@ -24,6 +24,7 @@ using WriteVTK using FillArrays using BlockArrays using LinearAlgebra +using ForwardDiff import Gridap.TensorValues: inner, outer, double_contraction, symmetric_part import LinearAlgebra: det, tr, cross, dot, ⋅, diag @@ -62,4 +63,6 @@ include("ODEs.jl") include("Adaptivity.jl") +include("Autodiff.jl") + end # module diff --git a/test/AutodiffTests.jl b/test/AutodiffTests.jl new file mode 100644 index 0000000..176a02e --- /dev/null +++ b/test/AutodiffTests.jl @@ -0,0 +1,45 @@ +module AtodiffTests + +using Test +using Gridap, Gridap.Algebra +using GridapDistributed +using PartitionedArrays + +function main(distribute,parts) + ranks = distribute(LinearIndices((prod(parts),))) + + domain = (0,4,0,4) + cells = (4,4) + model = CartesianDiscreteModel(ranks,parts,domain,cells) + + u((x,y)) = (x+y)^k + σ(∇u) = (1.0+∇u⋅∇u)*∇u + dσ(∇du,∇u) = (2*∇u⋅∇du)*∇u + (1.0+∇u⋅∇u)*∇du + f(x) = -divergence(y->σ(∇(u,y)),x) + + k = 1 + reffe = ReferenceFE(lagrangian,Float64,k) + V = TestFESpace(model,reffe,dirichlet_tags="boundary") + U = TrialFESpace(u,V) + + Ω = Triangulation(model) + dΩ = Measure(Ω,2*k) + r(u,v) = ∫( ∇(v)⋅(σ∘∇(u)) - v*f )dΩ + j(u,du,v) = ∫( ∇(v)⋅(dσ∘(∇(du),∇(u))) )dΩ + + op = FEOperator(r,j,U,V) + op_AD = FEOperator(r,U,V) + + uh = interpolate(1.0,U) + A = jacobian(op,uh) + A_AD = jacobian(op_AD,uh) + @test reduce(&,map(≈,partition(A),partition(A_AD))) + + g(v) = ∫(0.5*v⋅v)dΩ + dg(v) = ∫(uh⋅v)dΩ + b = assemble_vector(dg,U) + b_AD = assemble_vector(gradient(g,uh),U) + @test b ≈ b_AD +end + +end \ No newline at end of file diff --git a/test/PLaplacianTests.jl b/test/PLaplacianTests.jl index 2347dfc..a20a268 100644 --- a/test/PLaplacianTests.jl +++ b/test/PLaplacianTests.jl @@ -8,13 +8,12 @@ using Test using SparseArrays function main(distribute,parts) - main(distribute,parts,FullyAssembledRows(),SparseMatrixCSR{0,Float64,Int}) - main(distribute,parts,SubAssembledRows(),SparseMatrixCSC{Float64,Int}) + main(distribute,parts,FullyAssembledRows(),SparseMatrixCSR{0,Float64,Int},false) + main(distribute,parts,SubAssembledRows(),SparseMatrixCSC{Float64,Int},true) end -function main(distribute,parts,strategy,local_matrix_type) +function main(distribute,parts,strategy,local_matrix_type,autodiff) ranks = distribute(LinearIndices((prod(parts),))) - output = mkpath(joinpath(@__DIR__,"output")) domain = (0,4,0,4) cells = (4,4) @@ -22,7 +21,7 @@ function main(distribute,parts,strategy,local_matrix_type) k = 1 u((x,y)) = (x+y)^k - σ(∇u) =(1.0+∇u⋅∇u)*∇u + σ(∇u) = (1.0+∇u⋅∇u)*∇u dσ(∇du,∇u) = (2*∇u⋅∇du)*∇u + (1.0+∇u⋅∇u)*∇du f(x) = -divergence(y->σ(∇(u,y)),x) @@ -35,8 +34,12 @@ function main(distribute,parts,strategy,local_matrix_type) V = TestFESpace(model,reffe,dirichlet_tags="boundary") U = TrialFESpace(u,V) - assem=SparseMatrixAssembler(local_matrix_type,Vector{Float64},U,V,strategy) - op = FEOperator(r,j,U,V,assem) + assem = SparseMatrixAssembler(local_matrix_type,Vector{Float64},U,V,strategy) + if !autodiff + op = FEOperator(r,j,U,V,assem) + else + op = FEOperator(r,U,V,assem) + end uh = zero(U) b,A = residual_and_jacobian(op,uh) @@ -59,7 +62,6 @@ function main(distribute,parts,strategy,local_matrix_type) dΩo = Measure(Ωo,2*k) eh = u - uh @test sqrt(sum(∫( abs2(eh) )dΩo)) < 1.0e-9 - end end # module diff --git a/test/TestApp/src/TestApp.jl b/test/TestApp/src/TestApp.jl index d5c7f90..5a8390e 100644 --- a/test/TestApp/src/TestApp.jl +++ b/test/TestApp/src/TestApp.jl @@ -18,4 +18,5 @@ module TestApp include("../../AdaptivityMultiFieldTests.jl") include("../../BlockSparseMatrixAssemblersTests.jl") include("../../VisualizationTests.jl") + include("../../AutodiffTests.jl") end \ No newline at end of file diff --git a/test/mpi/runtests_np4_body.jl b/test/mpi/runtests_np4_body.jl index 5e484cf..59711fc 100644 --- a/test/mpi/runtests_np4_body.jl +++ b/test/mpi/runtests_np4_body.jl @@ -59,5 +59,8 @@ function all_tests(distribute,parts) PArrays.toc!(t,"Visualization") end + TestApp.AutodiffTests.main(distribute,parts) + PArrays.toc!(t,"Autodiff") + display(t) end diff --git a/test/sequential/AutodiffTests.jl b/test/sequential/AutodiffTests.jl new file mode 100644 index 0000000..3cbc5e8 --- /dev/null +++ b/test/sequential/AutodiffTests.jl @@ -0,0 +1,7 @@ +module AutodiffTestsSeq +using PartitionedArrays +include("../AutodiffTests.jl") +with_debug() do distribute + AutodiffTests.main(distribute,(2,2)) +end +end # module