Merge pull request #162 from Cthonios/chchchchanges

block arrays ext, initial stab at matrix free assembler and adding ex…

Author: cmhamel

Project.toml

@@ -18,12 +18,14 @@ TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Exodus = "f57ae99e-f805-4780-bdca-96e224be1e5a"
 
 [extensions]
 FiniteElementContainersAMDGPUExt = ["Adapt", "AMDGPU"]
 FiniteElementContainersAdaptExt = "Adapt"
+FiniteElementContainersBlockArraysExt = "BlockArrays"
 FiniteElementContainersCUDAExt = ["Adapt", "CUDA"]
 FiniteElementContainersExodusExt = "Exodus"
 
@@ -32,6 +34,7 @@ AMDGPU = "2"
 Adapt = "4"
 Aqua = "0.8"
 Atomix = "1"
+BlockArrays = "1"
 CUDA = "5"
 DocStringExtensions = "0.9"
 Exodus = "0.14"

ext/FiniteElementContainersBlockArraysExt.jl

@@ -0,0 +1,11 @@
+module FiniteElementContainersBlockArraysExt
+
+using BlockArrays
+using FiniteElementContainers
+
+function FiniteElementContainers.create_unknowns(assemblers::NamedTuple)
+    Uus = map(create_unknowns, collect(values(assemblers)))
+    return mortar(Uus)
+end
+
+end # module

src/DofManagers.jl

@@ -62,6 +62,36 @@ function create_unknowns(dof::DofManager)
     return KA.zeros(backend, Float64, length(dof.unknown_dofs))
 end
 
+# COV_EXCL_START
+KA.@kernel function _extract_field_unknowns_kernel!(Uu::V, dof::DofManager, U::AbstractField) where V <: AbstractVector{<:Number}
+    N = KA.@index(Global)
+    @inbounds Uu[N] = U[dof.unknown_dofs[N]]
+end
+# COV_EXCL_STOP
+
+function _extract_field_unknowns!(Uu::V, dof::DofManager, U::AbstractField, backend::KA.Backend) where V <: AbstractVector{<:Number}
+    kernel! = _extract_field_unknowns_kernel!(backend)
+    kernel!(Uu, dof, U, ndrange = length(Uu))
+    return nothing
+end
+
+function _extract_field_unknowns!(Uu::V, dof::DofManager, U::AbstractField, ::KA.CPU) where V <: AbstractVector{<:Number}
+    @views Uu .= U[dof.unknown_dofs]
+    return nothing
+end
+
+function extract_field_unknowns!(
+    Uu::V,
+    dof::DofManager,
+    U::AbstractField
+) where V <: AbstractVector{<:Number}
+    backend = KA.get_backend(dof)
+    @assert KA.get_backend(U) == backend
+    @assert KA.get_backend(Uu) == backend
+    _extract_field_unknowns!(Uu, dof, U, backend)
+    return nothing
+end
+
 function function_space(dof::DofManager)
     return dof.var.fspace
 end
@@ -84,20 +114,20 @@ function update_dofs!(dof::DofManager, dirichlet_dofs)
 end
 
 # COV_EXCL_START
-KA.@kernel function _update_field_unknowns_kernel!(U::H1Field, dof::DofManager, Uu::T) where T <: AbstractArray{<:Number, 1}
+KA.@kernel function _update_field_unknowns_kernel!(U::AbstractField, dof::DofManager, Uu::V) where V <: AbstractVector{<:Number}
     N = KA.@index(Global)
     @inbounds U[dof.unknown_dofs[N]] = Uu[N]
 end
 # COV_EXCL_STOP
 
-function _update_field_unknowns!(U::H1Field, dof::DofManager, Uu::T, backend::KA.Backend) where T <: AbstractArray{<:Number, 1}
+function _update_field_unknowns!(U::AbstractField, dof::DofManager, Uu::T, backend::KA.Backend) where T <: AbstractVector{<:Number}
     kernel! = _update_field_unknowns_kernel!(backend)
     kernel!(U, dof, Uu, ndrange = length(Uu))
     return nothing
 end
 
 # Need a seperate CPU method since CPU is basically busted in KA
-function _update_field_unknowns!(U::H1Field, dof::DofManager, Uu::T, ::KA.CPU) where T <: AbstractArray{<:Number, 1}
+function _update_field_unknowns!(U::AbstractField, dof::DofManager, Uu::T, ::KA.CPU) where T <: AbstractVector{<:Number}
     U[dof.unknown_dofs] .= Uu
     return nothing
 end

src/assemblers/Assemblers.jl

@@ -172,7 +172,12 @@ end
 $(TYPEDSIGNATURES)
 """
 function hvp(assembler::AbstractAssembler)
-  return _hvp(assembler, KA.get_backend(assembler))
+  extract_field_unknowns!(
+    asm.stiffness_action_unknowns,
+    asm.dof,
+    asm.stiffness_action_storage
+  )
+  return nothing
 end
 
 """
@@ -186,7 +191,12 @@ end
 $(TYPEDSIGNATURES)
 """
 function residual(asm::AbstractAssembler)
-  return _residual(asm, KA.get_backend(asm))
+  extract_field_unknowns!(
+    asm.residual_unknowns, 
+    asm.dof, 
+    asm.residual_storage
+  )
+  return asm.residual_unknowns
 end
 
 """
@@ -200,6 +210,7 @@ end
 include("SparsityPattern.jl")
 
 # types
+include("MatrixFreeAssembler.jl")
 include("SparseMatrixAssembler.jl")
 
 # methods

src/assemblers/MatrixAction.jl

@@ -80,18 +80,6 @@ function _assemble_block_matrix_action!(
   end
 end
 
-# TODO hardcoded to H1 fields right now.
-"""
-$(TYPEDSIGNATURES)
-"""
-function _hvp(asm::AbstractAssembler, ::KA.CPU)
-  # for n in axes(asm.residual_unknowns, 1)
-  #   asm.residual_unknowns[n] = asm.residual_storage[asm.dof.H1_unknown_dofs[n]]
-  # end
-  @views asm.stiffness_action_unknowns .= asm.stiffness_action_storage[asm.dof.H1_unknown_dofs]
-  return asm.stiffness_action_unknowns
-end
-
 # GPU implementation
 
 """
@@ -187,26 +175,3 @@ function _assemble_block_matrix_action!(
   )
   return nothing
 end
-
-# """
-# $(TYPEDSIGNATURES)
-# """
-# KA.@kernel function _extract_residual_unknowns!(Ru, unknown_dofs, R)
-#   N = KA.@index(Global)
-#   Ru[N] = R[unknown_dofs[N]]
-# end
-
-"""
-$(TYPEDSIGNATURES)
-TODO note will only work for H1 spaces right now.
-  TODO move me to Assemblers.jl
-"""
-function _hvp(asm::AbstractAssembler, backend::KA.Backend)
-  kernel! = _extract_residual_unknowns!(backend)
-  kernel!(asm.stiffness_action_unknowns, 
-          asm.dof.H1_unknown_dofs, 
-          asm.stiffness_action_storage, 
-          ndrange=length(asm.dof.H1_unknown_dofs))
-  KA.synchronize(KA.get_backend(asm))
-  return asm.stiffness_action_unknowns
-end 

src/assemblers/MatrixFreeAssembler.jl

@@ -0,0 +1,57 @@
+struct MatrixFreeAssembler{
+    Dof <: DofManager,
+    Storage1 <: AbstractField,
+    Storage2 <: AbstractArray{<:Number, 1}
+} <: AbstractAssembler{Dof}
+    dof::Dof
+    residual_storage::Storage1
+    residual_unknowns::Storage2
+    stiffness_action_storage::Storage1
+    stiffness_action_unknowns::Storage2
+end
+
+function MatrixFreeAssembler(dof::DofManager)
+
+    residual_storage = create_field(dof)
+    residual_unknowns = create_unknowns(dof)
+    stiffness_action_storage = create_field(dof)
+    stiffness_action_unknowns = create_unknowns(dof)
+    return MatrixFreeAssembler(
+        dof,
+        residual_storage, residual_unknowns,
+        stiffness_action_storage, stiffness_action_unknowns
+    )
+end
+
+function MatrixFreeAssembler(u::AbstractFunction)
+    dof = DofManager(u)
+    return MatrixFreeAssembler(dof)
+end
+
+function update_dofs!(assembler::MatrixFreeAssembler, dirichlet_bcs)
+    # collect dbcs
+    if length(dirichlet_bcs) > 0
+        dirichlet_dofs = mapreduce(x -> x.bookkeeping.dofs, vcat, dirichlet_bcs)
+        dirichlet_dofs = unique(sort(dirichlet_dofs))
+    else
+        dirichlet_dofs = Vector{Int}(undef, 0)
+    end
+
+    # update dof manager
+    update_dofs!(assembler.dof, dirichlet_dofs)
+
+    # update cached arrays in assembler
+    resize!(assembler.residual_unknowns, length(assembler.dof.unknown_dofs))
+    resize!(assembler.stiffness_action_unknowns, length(assembler.dof.unknown_dofs))
+
+    return nothing
+end
+
+Base.eltype(asm::MatrixFreeAssembler) = eltype(asm.residual_storage)
+Base.size(asm::MatrixFreeAssembler) = (length(asm.residual_unknowns), length(asm.residual_unknowns))
+
+function LinearAlgebra.mul!(y::T, asm::MatrixFreeAssembler, Uu::T, p) where T
+    assemble!(asm.stiffness_action_storage, hvp, Uu, p)
+    copyto!(y, hvp(asm))
+    return nothing
+end

src/assemblers/Vector.jl

@@ -81,15 +81,6 @@ function _assemble_block_vector!(
   end
 end
 
-# TODO hardcoded to H1 fields right now.
-"""
-$(TYPEDSIGNATURES)
-"""
-function _residual(asm::AbstractAssembler, ::KA.CPU)\
-  @views asm.residual_unknowns .= asm.residual_storage[asm.dof.unknown_dofs]
-  return asm.residual_unknowns
-end
-
 # GPU implementation
 
 """
@@ -184,28 +175,3 @@ function _assemble_block_vector!(
   )
   return nothing
 end
-
-"""
-$(TYPEDSIGNATURES)
-"""
-# COV_EXCL_START
-KA.@kernel function _extract_residual_unknowns!(Ru, unknown_dofs, R)
-  N = KA.@index(Global)
-  Ru[N] = R[unknown_dofs[N]]
-end
-# COV_EXCL_STOP
-
-"""
-$(TYPEDSIGNATURES)
-"""
-function _residual(asm::AbstractAssembler, backend::KA.Backend)
-  kernel! = _extract_residual_unknowns!(backend)
-  kernel!(
-    asm.residual_unknowns, 
-    asm.dof.unknown_dofs, 
-    asm.residual_storage, 
-    ndrange=length(asm.dof.unknown_dofs)
-  )
-  KA.synchronize(KA.get_backend(asm))
-  return asm.residual_unknowns
-end