Merge pull request #80 from Cthonios/assemblers-gpu

Assemblers gpu

Author: cmhamel

Project.toml

@@ -23,7 +23,7 @@ Exodus = "f57ae99e-f805-4780-bdca-96e224be1e5a"
 
 [extensions]
 FiniteElementContainersAdaptExt = "Adapt"
-FiniteElementContainersCUDAExt = "CUDA"
+FiniteElementContainersCUDAExt = ["Adapt", "CUDA"]
 FiniteElementContainersExodusExt = "Exodus"
 
 [compat]
@@ -36,6 +36,7 @@ DocStringExtensions = "0.9"
 Exodus = "0.13"
 IterativeSolvers = "0.9"
 JET = "0.9"
+Krylov = "0.9"
 KernelAbstractions = "0.9"
 LinearAlgebra = "1"
 Parameters = "0.12"
@@ -55,9 +56,10 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Exodus = "f57ae99e-f805-4780-bdca-96e224be1e5a"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestSetExtensions = "98d24dd4-01ad-11ea-1b02-c9a08f80db04"
 
 [targets]
-test = ["Aqua", "CUDA", "Exodus", "IterativeSolvers", "JET", "Parameters", "Test", "TestSetExtensions"]
+test = ["Aqua", "CUDA", "Exodus", "IterativeSolvers", "JET", "Krylov", "Parameters", "Test", "TestSetExtensions"]

ext/FiniteElementContainersAdaptExt.jl

@@ -3,16 +3,21 @@ module FiniteElementContainersAdaptExt
 using Adapt
 using FiniteElementContainers
 
+FiniteElementContainersAdaptExt.cpu(x) = Adapt.adapt_structure(Array, x)
+
 # Assemblers
 function Adapt.adapt_structure(to, asm::SparseMatrixAssembler)
   dof = Adapt.adapt_structure(to, asm.dof)
   pattern = Adapt.adapt_structure(to, asm.pattern)
   constraint_storage = Adapt.adapt_structure(to, asm.constraint_storage)
   residual_storage = Adapt.adapt_structure(to, asm.residual_storage)
+  residual_unknowns = Adapt.adapt_structure(to, asm.residual_unknowns)
   stiffness_storage = Adapt.adapt_structure(to, asm.stiffness_storage)
   return SparseMatrixAssembler(
     dof, pattern, 
-    constraint_storage, residual_storage, stiffness_storage
+    constraint_storage, 
+    residual_storage, residual_unknowns,
+    stiffness_storage
   )
 end
 

ext/FiniteElementContainersCUDAExt.jl

@@ -1,16 +1,19 @@
 module FiniteElementContainersCUDAExt
 
+using Adapt
 using CUDA
 using FiniteElementContainers
 using KernelAbstractions
 
+FiniteElementContainers.gpu(x) = Adapt.adapt_structure(CuArray, x)
+
 # this method need to have the assembler initialized first
 # if the stored values in asm.pattern.cscnzval or zero
 # CUDA will error out
 function CUDA.CUSPARSE.CuSparseMatrixCSC(asm::SparseMatrixAssembler)
   @assert typeof(get_backend(asm)) <: CUDABackend "Assembler is not on a CUDA device"
-  @assert length(asm.pattern.cscnzval) > 0 "Need to assemble the assembler once"
-  @assert all(x -> x != zero(eltype(asm.pattern.cscnzval)), asm.pattern.cscnzval) "Need to assemble the assembler once"
+  @assert length(asm.pattern.cscnzval) > 0 "Need to assemble the assembler once with SparseArrays.sparse!(assembler)"
+  @assert all(x -> x != zero(eltype(asm.pattern.cscnzval)), asm.pattern.cscnzval) "Need to assemble the assembler once with SparseArrays.sparse!(assembler)"
   n_dofs = FiniteElementContainers.num_unknowns(asm.dof)
   return CUDA.CUSPARSE.CuSparseMatrixCSC(
     asm.pattern.csccolptr,
@@ -20,6 +23,10 @@ function CUDA.CUSPARSE.CuSparseMatrixCSC(asm::SparseMatrixAssembler)
   )
 end
 
+function FiniteElementContainers._stiffness(asm::SparseMatrixAssembler, ::Backend)
+  return CUDA.CUSPARSE.CuSparseMatrixCSC(asm)
+end
+
 # this one isn't quite right
 # function CUDA.CUSPARSE.CuSparseMatrixCSR(asm::SparseMatrixAssembler)
 #   n_dofs = FiniteElementContainers.num_unknowns(asm.dof)

src/Assemblers.jl

@@ -256,6 +256,135 @@ function assemble!(
   return nothing
 end
 
+# wrapper
+function assemble!(
+  assembler::SparseMatrixAssembler,
+  func::F,
+  U::T,
+  sym::Symbol
+) where {F, T}
+  assemble!(assembler, func, U, Val{sym}())
+end
+
+KA.@kernel function _assemble_residual_kernel!(asm, func, U, X, conns, ref_fe, block_id)
+  E = KA.@index(Global)
+
+  # some helpers
+  NDim = size(X, 1)
+  ND, NN = num_dofs_per_node(asm.dof), num_nodes(asm.dof)
+  NNPE = ReferenceFiniteElements.num_vertices(ref_fe)
+  # NDim = ReferenceFiniteElements.num_dimensions(ref_fe)
+  NxNDof = NNPE * ND
+  # IDs = reshape(1:length(asm.dof), ND, size(U, 2))
+
+  # extract element level stuff
+  u_el = @views SMatrix{ND, NNPE, Float64, NxNDof}(U[:, conns[:, E]])
+  x_el = @views SMatrix{NDim, NNPE, Float64, NDim * NNPE}(X[:, conns[:, E]])
+  # dof_conns = reshape(IDs[:, conns], ND * length(conns))
+
+  # assemble element level residual
+  R_el = zeros(SVector{NxNDof, Float64})
+  for q in 1:num_quadrature_points(ref_fe)
+    interps = MappedInterpolants(ref_fe.cell_interps.vals[q], x_el)
+
+    # dynamic method invocation isn't working here. 
+    # R_q = func(interps, u_el)
+    # R_el = R_el + R_q
+
+    X_q, N, ∇N_X, JxW = interps.X_q, interps.N, interps.∇N_X, interps.JxW
+    ∇u_q = u_el * ∇N_X
+    f = 2. * π^2 * sin(2π * X_q[1]) * sin(4π * X_q[2])
+    R_q = ∇u_q * ∇N_X' - N' * f
+    R_el = R_el + JxW * R_q[:]
+  end
+
+  # now assemble
+  # TODO won't work for problems with more than one dof...
+  # need to create dof_conns
+  for i in axes(conns, 1)
+    Atomix.@atomic asm.residual_storage.vals[conns[i]] += R_el[i]
+  end
+end
+
+KA.@kernel function _assemble_stiffness_kernel!(asm, func, U, X, conns, ref_fe, block_id)
+  E = KA.@index(Global)
+
+  # some helpers
+  NDim = size(X, 1)
+  ND, NN = num_dofs_per_node(asm.dof), num_nodes(asm.dof)
+  NNPE = ReferenceFiniteElements.num_vertices(ref_fe)
+  # NDim = ReferenceFiniteElements.num_dimensions(ref_fe)
+  NxNDof = NNPE * ND
+
+  # extract element level stuff
+  u_el = @views SMatrix{ND, NNPE, Float64, NxNDof}(U[:, conns[:, E]])
+  x_el = @views SMatrix{NDim, NNPE, Float64, NDim * NNPE}(X[:, conns[:, E]])
+
+  # assemble element level residual
+  K_el = zeros(SMatrix{NxNDof, NxNDof, Float64, NxNDof * NxNDof})
+  for q in 1:num_quadrature_points(ref_fe)
+    interps = MappedInterpolants(ref_fe.cell_interps.vals[q], x_el)
+
+    # dynamic method invocation isn't working here. 
+    # R_q = func(interps, u_el)
+    # R_el = R_el + R_q
+
+    X_q, N, ∇N_X, JxW = interps.X_q, interps.N, interps.∇N_X, interps.JxW
+    K_q = ∇N_X * ∇N_X'
+    K_el = K_el + JxW * K_q
+  end
+
+  # now assemble
+  # TODO won't work for problems with more than one dof...
+  # need to create dof_conns
+  # start_id = (block_id - 1) * asm.pattern.block_sizes[block_id] + 
+  #            (el_id - 1) * asm.pattern.block_offsets[block_id] + 1
+  # end_id = start_id + asm.pattern.block_offsets[block_id] - 1
+  # ids = start_id:end_id
+  # for id in ids
+  #   Atomix.@atomic asm.stiffness_storage[id] += K_el[:]
+  # end
+end
+
+# newer method try to dispatch based on field type
+function assemble!(
+  assembler::SparseMatrixAssembler,
+  residual_func::F,
+  U::T,
+  ::Val{:residual}
+) where {F <: Function, T <: H1Field}
+  # TODO hack for now, only grabbing first var fspace
+  # fspace = getproperty(assembler.dof.H1_vars, :H1_var_1)
+  fspace = assembler.dof.H1_vars[1].fspace
+  for (block_id, conns) in enumerate(values(fspace.elem_conns))
+    ref_fe = values(fspace.ref_fes)[block_id]
+    # backend = KA.get_backend(assembler)
+    # TODO eventually dispatch on CPU vs. GPU
+    kernel! = _assemble_residual_kernel!(KA.get_backend(assembler))
+    kernel!(assembler, residual_func, U, fspace.coords, conns, ref_fe, block_id, ndrange = size(conns, 2))
+  end
+end
+
+function assemble!(
+  assembler::SparseMatrixAssembler,
+  tangent_func::F,
+  U::T,
+  ::Val{:stiffness}
+) where {F <: Function, T <: H1Field}
+  # TODO hack for now, only grabbing first var fspace
+  # fspace = getproperty(assembler.dof.H1_vars, :H1_var_1)
+  fspace = assembler.dof.H1_vars[1].fspace
+  for (block_id, conns) in enumerate(values(fspace.elem_conns))
+    ref_fe = values(fspace.ref_fes)[block_id]
+    # backend = KA.get_backend(assembler)
+    # TODO eventually dispatch on CPU vs. GPU
+
+    # TODO going to need to query the start/end ids for the block assembler
+    kernel! = _assemble_stiffness_kernel!(KA.get_backend(assembler))
+    kernel!(assembler, tangent_func, U, fspace.coords, conns, ref_fe, block_id, ndrange = size(conns, 2))
+  end
+end
+
 # TODO hardcoded for H1
 function assemble!(
   R,
@@ -265,6 +394,7 @@ function assemble!(
 ) where {F <: Function, T <: AbstractArray}
 
   R .= zero(eltype(R))
+  # fill(R.vals, zero(eltype(R)))
 
   vars = assembler.dof.H1_vars
 
@@ -274,42 +404,9 @@ function assemble!(
 
   fspace = vars[1].fspace
 
-  # NDim = size(fspace.coords, 1)
-  # ND, NN = num_dofs_per_node(assembler.dof), num_nodes(assembler.dof)
-  # ids = reshape(1:length(assembler.dof), ND, NN)
-
   for (block_id, conns) in enumerate(values(fspace.elem_conns))
     ref_fe = values(fspace.ref_fes)[block_id]
-
     assemble!(R, assembler, residual_func, U, fspace.coords, conns, ref_fe)
-    # NNPE = ReferenceFiniteElements.num_vertices(ref_fe)
-    # NxNDof = NNPE * ND
-    # dof_conns = @views reshape(ids[:, conns], ND * size(conns, 1), size(conns, 2))
-
-    # for e in 1:size(conns, 2)
-    # AK.foraxes(conns, 2) do e
-    #   u_el = @views SMatrix{ND, NNPE, Float64, NxNDof}(U[dof_conns[:, e]])
-    #   x_el = @views SMatrix{NDim, NNPE, Float64, NDim * NNPE}(fspace.coords[:, conns[:, e]])
-    #   # K_el = zeros(SMatrix{NxNDof, NxNDof, Float64, NxNDof * NxNDof})
-    #   # R_el = zeros(SMatrix{ND, NNPE, Float64, NxNDof})
-    #   R_el = zeros(SVector{NxNDof, Float64})
-
-    #   for q in 1:num_quadrature_points(ref_fe)
-    #     N = ReferenceFiniteElements.shape_function_value(ref_fe, q)
-    #     ∇N_ξ = ReferenceFiniteElements.shape_function_gradient(ref_fe, q)
-    #     ∇N_X = map_shape_function_gradients(x_el, ∇N_ξ)
-    #     JxW  = volume(x_el, ∇N_ξ) * quadrature_weight(ref_fe, q)
-    #     x_q = x_el * N
-    #     interps = Interpolants(x_q, N, ∇N_X, JxW)
-        
-    #     R_q = residual_func(interps, u_el)
-    #     R_el = R_el + JxW * R_q
-    #   end
-
-    #   # assemble!(assembler, R_el, e, block_id)
-    #   # Note this method is in old stuff TODO
-    #   @views assemble!(R, R_el, dof_conns[:, e])
-    # end
   end
 end
 
@@ -396,6 +493,7 @@ function assemble!(
   end
 end
 
+create_bcs(asm::SparseMatrixAssembler, type) = create_bcs(asm.dof, type)
 create_field(asm::SparseMatrixAssembler, type) = create_field(asm.dof, type)
 create_unknowns(asm::SparseMatrixAssembler) = create_unknowns(asm.dof)
 

src/DofManagers.jl

@@ -116,6 +116,10 @@ function _dof_manager_sym_name(u::VectorFunction)
   return Symbol(split(String(names(u)[1]), ['_'])[1])
 end
 
+function _dof_manager_vars(dof::DofManager, ::Type{<:H1Field})
+  return dof.H1_vars
+end
+
 function _filter_field_type(vars, T)
   vars = filter(x -> isa(x.fspace.coords, T), vars)
   syms = map(_dof_manager_sym_name, vars)
@@ -167,6 +171,14 @@ Base.length(dof::DofManager) = length(dof.H1_bc_dofs) + length(dof.H1_unknown_do
 
 KA.get_backend(dof::DofManager) = KA.get_backend(dof.H1_unknown_dofs)
 
+"""
+$(TYPEDSIGNATURES)
+"""
+function create_bcs(dof::DofManager, ::Type{H1Field})
+  backend = KA.get_backend(dof.H1_bc_dofs)
+  return KA.zeros(backend, Float64, length(dof.H1_bc_dofs))
+end
+
 """
 $(TYPEDSIGNATURES)
 """
@@ -361,6 +373,33 @@ function update_field_unknowns!(U::H1Field, dof::DofManager, Uu::T) where T <: A
   return nothing
 end
 
+KA.@kernel function _update_field_unknowns_kernel_plus!(U::H1Field, dof::DofManager, Uu::T) where T <: AbstractArray{<:Number, 1}
+  N = KA.@index(Global)
+  @inbounds U[dof.H1_unknown_dofs[N]] += Uu[N]
+end
+
+function _update_field_unknowns!(U::H1Field, dof::DofManager, Uu::T, ::typeof(+), backend::KA.Backend) where T <: AbstractArray{<:Number, 1}
+  kernel! = _update_field_unknowns_kernel_plus!(backend)
+  kernel!(U, dof, Uu, ndrange = length(Uu))
+  return nothing
+end
+
+function _update_field_unknowns!(U::H1Field, dof::DofManager, Uu::T, ::typeof(+), ::KA.CPU) where T <: AbstractArray{<:Number, 1}
+  for (n, d) in enumerate(dof.H1_unknown_dofs)
+    U[d] += Uu[n]
+  end
+  return nothing
+end
+
+"""
+$(TYPEDSIGNATURES)
+Does a simple addition on CPUs. On GPUs it uses a ```KernelAbstractions``` kernel
+"""
+function update_field_unknowns!(U::H1Field, dof::DofManager, Uu::T, ::typeof(+)) where T <: AbstractArray{<:Number, 1}
+  _update_field_unknowns!(U, dof, Uu, +, KA.get_backend(dof))
+  return nothing
+end
+
 """
 $(TYPEDSIGNATURES)
 """

src/FiniteElementContainers.jl

@@ -1,8 +1,14 @@
 module FiniteElementContainers
 
+# general
+export cpu
+export gpu
+
 # Assemblers
 export SparseMatrixAssembler
 export assemble!
+export residual
+export stiffness
 
 # BCs
 export DirichletBC
@@ -41,6 +47,7 @@ export connectivity
 
 # DofManager
 export DofManager
+export create_bcs
 export create_field
 export create_unknowns
 export update_dofs!
@@ -87,6 +94,9 @@ using Tensors
 
 abstract type FEMContainer end
 
+function cpu end
+function gpu end
+
 # basic stuff
 include("fields/Fields.jl")
 include("Meshes.jl")
@@ -100,6 +110,9 @@ include("Functions.jl")
 include("DofManagers.jl")
 include("bcs/BoundaryConditions.jl")
 include("formulations/Formulations.jl")
-include("Assemblers.jl")
+# include("Assemblers.jl")
+include("assemblers/Assemblers.jl")
+
+include("physics/Physics.jl")
 
 end # module