Merge pull request #41 from Cthonios/atomic-assemble

Atomic assemble

Project.toml

@@ -4,6 +4,7 @@ authors = ["Craig M. Hamel <cmhamel32@gmail.com> and contributors"]
 version = "0.4.3"
 
 [deps]
+Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ReferenceFiniteElements = "6dc62d09-f8eb-43fd-9672-074e490a997f"
@@ -25,6 +26,7 @@ FiniteElementContainersKernelAbstractionsExt = "KernelAbstractions"
 [compat]
 Adapt = "3"
 Aqua = "0.8"
+Atomix = "0.1.0"
 DocStringExtensions = "0.9"
 Exodus = "0.12"
 IterativeSolvers = "0.9"

src/FiniteElementContainers.jl

@@ -50,10 +50,12 @@ export num_q_points
 # Assemblers
 export Assembler
 export DynamicAssembler
+export MatrixFreeAssembler
 export StaticAssembler
 export assemble!
 
 # dependencies
+using Atomix
 using DocStringExtensions
 using LinearAlgebra
 using ReferenceFiniteElements

src/assemblers/Assemblers.jl

@@ -31,6 +31,23 @@ function assemble!(
   end
 end
 
+"""
+$(TYPEDSIGNATURES)
+assembly method for just a residual vector
+
+TODO need to add an Atomix lock here
+TODO add block_id to fspace or something like that
+"""
+function assemble_atomic!(
+  assembler::Assembler,
+  R_el::V1, conn::V2
+) where {V1 <: AbstractVector{<:Number}, V2 <: AbstractVector{<:Integer}}
+
+  for i in axes(conn, 1)
+    Atomix.@atomic assembler.residuals.vals[conn[i]] += R_el[i]
+  end
+end
+
 """
 $(TYPEDSIGNATURES)
 method that assumes first dof
@@ -80,4 +97,5 @@ end
 
 # implementations
 include("DynamicAssemblers.jl")
+include("MatrixFreeAssembler.jl")
 include("StaticAssembler.jl")

src/assemblers/MatrixFreeAssembler.jl

@@ -0,0 +1,40 @@
+
+struct MatrixFreeAssembler{Rtype, Itype} <: Assembler{Rtype, Itype}
+end
+
+"""
+assembly method for matrix free assembler
+when functions for the element level residual and tangent_func
+are provided
+"""
+function assemble!(
+  assembler::MatrixFreeAssembler,
+  dof::DofManager,
+  fspace::FunctionSpace,
+  X, U, V,
+  residual_func, tangent_func
+)
+
+  NDof = num_dofs_per_node(dof)
+  N    = num_nodes_per_element(fspace)
+  NxNDof = N * NDof
+
+  for e in 1:num_elements(fspace)
+    U_el = element_level_fields(fspace, U, e)
+    V_el = element_level_fields(fspace, V, e)
+    R_el = zeros(SVector{NxNDof, Float64})
+    Hv_el = zeros(SVector{NxNDof, Float64})
+
+    # quadrature loop
+    for q in 1:num_q_points(fspace)
+      fspace_values = getindex(fspace, X, q, e)
+      R_el = R_el + residual_func(fspace_values, U_el)
+      Hv_el = Hv_el + dot(tangent_func(fspace_values, U_el), V_el)
+    end
+
+    # assemble residual using connectivity here
+    conn = dof_connectivity(fspace, e)
+    assemble!(assembler, R_el, conn)
+    assemble!(assembler, Hv_el, conn)
+  end
+end

src/assemblers/StaticAssembler.jl

@@ -147,6 +147,28 @@ function assemble!(
   assembler.stiffnesses[ids] = K_el
 end
 
+"""
+$(TYPEDSIGNATURES)
+assembly for stiffness matrix
+"""
+function assemble_atomic!(
+  assembler::StaticAssembler,
+  K_el::M, block_id::Int, el_id::Int
+) where M <: AbstractMatrix
+
+  # first get mapping from block and element id to ids in assembler.stiffnesses
+  start_id = (block_id - 1) * assembler.block_sizes[block_id] + 
+             (el_id - 1) * assembler.block_offsets[block_id] + 1
+  end_id = start_id + assembler.block_offsets[block_id] - 1
+  ids = start_id:end_id
+
+  # now assemble into stiffnesses
+  # Atomix.@atomic assembler.stiffnesses[ids] = K_el
+  for (n, id) in enumerate(ids)
+    Atomix.@atomic assembler.stiffnesses[id] = K_el[n]
+  end
+end
+
 """
 $(TYPEDSIGNATURES)
 Simple method for assembling in serial