From 422c06d5e727c9778b3a882a8dc3c748c9fd1e29 Mon Sep 17 00:00:00 2001 From: "Craig M. Hamel" Date: Fri, 28 Nov 2025 18:00:08 -0500 Subject: [PATCH] cleaning up assemblers typing among other things. --- src/assemblers/Assemblers.jl | 29 +---- src/assemblers/Matrix.jl | 144 ++++++++++++------------ src/assemblers/MatrixAction.jl | 116 +++++++++---------- src/assemblers/NeumannBC.jl | 54 +++------ src/assemblers/QuadratureQuantity.jl | 113 +++++++++---------- src/assemblers/SparseMatrixAssembler.jl | 25 +--- src/assemblers/SparsityPatterns.jl | 43 +++++-- src/assemblers/Vector.jl | 115 +++++++++---------- 8 files changed, 293 insertions(+), 346 deletions(-) diff --git a/src/assemblers/Assemblers.jl b/src/assemblers/Assemblers.jl index 96f982d..23487e1 100644 --- a/src/assemblers/Assemblers.jl +++ b/src/assemblers/Assemblers.jl @@ -75,7 +75,7 @@ Assembly method for an H1Field, e.g. internal force Called on a single element e for a given block b where local_val has already been constructed from quadrature contributions. """ -function _assemble_element!(global_val::H1Field, local_val, conn, e, b) +function _assemble_element!(global_val::H1Field, local_val, conn) n_dofs = size(global_val, 1) for i in axes(conn, 1) for d in 1:n_dofs @@ -88,33 +88,6 @@ function _assemble_element!(global_val::H1Field, local_val, conn, e, b) return nothing end -# function _check_backends(assembler, U, X, state_old, state_new, conns) -# backend = KA.get_backend(assembler) -# # TODO add get_backend method of ref_fe -# @assert backend == KA.get_backend(U) -# @assert backend == KA.get_backend(X) -# @assert backend == KA.get_backend(conns) -# @assert backend == KA.get_backend(state_old) -# @assert backend == KA.get_backend(state_new) -# # props will be complicated... -# # TODO -# return backend -# end - -# function _check_backends(assembler, U, V, X, state_old, state_new, conns) -# backend = KA.get_backend(assembler) -# # TODO add get_backend method of ref_fe -# @assert backend == KA.get_backend(U) -# @assert backend == KA.get_backend(V) -# @assert backend == KA.get_backend(X) -# @assert backend == KA.get_backend(conns) -# @assert backend == KA.get_backend(state_old) -# @assert backend == KA.get_backend(state_new) -# # props will be complicated... -# # TODO -# return backend -# end - """ $(TYPEDSIGNATURES) """ diff --git a/src/assemblers/Matrix.jl b/src/assemblers/Matrix.jl index 5cff1c3..1ac9fab 100644 --- a/src/assemblers/Matrix.jl +++ b/src/assemblers/Matrix.jl @@ -22,32 +22,36 @@ Note this is hard coded to storing the assembled sparse matrix in the stiffness_storage field of assembler. """ function assemble_matrix!( - # assembler, func::F, Uu, p, ::Type{H1Field}; - # storage_sym=Val{:stiffness_storage}() storage, pattern, dof, func::F, Uu, p ) where F <: Function - # storage = getfield(assembler, storage_sym) - # storage = _get_storage(assembler, storage_sym) fill!(storage, zero(eltype(storage))) fspace = function_space(dof) t = current_time(p.times) dt = time_step(p.times) _update_for_assembly!(p, dof, Uu) - for (b, (conns, block_physics, state_old, state_new, props)) in enumerate(zip( + for ( + conns, block_start_index, block_el_level_size, + block_physics, ref_fe, + state_old, state_new, props + ) in zip( values(fspace.elem_conns), - values(p.physics), + values(pattern.block_start_indices), + values(pattern.block_el_level_sizes), + values(p.physics), values(fspace.ref_fes), values(p.state_old), values(p.state_new), - values(p.properties) - )) - ref_fe = values(fspace.ref_fes)[b] + values(p.properties), + ) # TODO re-enable back-end checking - # backend = _check_backends(assembler, p.h1_field, p.h1_coords, state_old, state_new, conns) backend = KA.get_backend(p.h1_field) _assemble_block_matrix!( - storage, pattern, block_physics, ref_fe, - p.h1_field, p.h1_field_old, p.h1_coords, state_old, state_new, props, t, dt, - conns, b, func, - backend + backend, + storage, + conns, block_start_index, block_el_level_size, + func, + block_physics, ref_fe, + p.h1_coords, t, dt, + p.h1_field, p.h1_field_old, + state_old, state_new, props ) end end @@ -62,22 +66,18 @@ with no threading. TODO add state variables and physics properties """ function _assemble_block_matrix!( - field::F1, pattern::Patt, physics::Phys, ref_fe::R, - U::F2, U_old::F2, X::F3, state_old::S, state_new::S, props::P, t::T, dt::T, - conns::C, block_id::Int, func::Func, ::KA.CPU + ::KA.CPU, + field::AbstractArray{<:Number, 1}, + conns::Connectivity, block_start_index::Int, block_el_level_size::Int, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, dt::T, + U::Solution, U_old::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - F1 <: AbstractArray{<:Number, 1}, - F2 <: AbstractField, - F3 <: AbstractField, - Func <: Function, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Patt <: SparseMatrixPattern, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } for e in axes(conns, 2) @@ -90,14 +90,13 @@ function _assemble_block_matrix!( for q in 1:num_quadrature_points(ref_fe) interps = _cell_interpolants(ref_fe, q) state_old_q = _quadrature_level_state(state_old, q, e) - # K_q, state_new_q = func(physics, interps, u_el, x_el, state_old_q, props_el, t, dt) K_q, state_new_q = func(physics, interps, x_el, t, dt, u_el, u_el_old, state_old_q, props_el) K_el = K_el + K_q for s in 1:length(state_old) state_new[s, q, e] = state_new_q[s] end end - @views _assemble_element!(pattern, field, K_el, e, block_id) + _assemble_element!(field, K_el, e, block_start_index, block_el_level_size) end return nothing end @@ -106,22 +105,17 @@ end # GPU implementation # COV_EXCL_START KA.@kernel function _assemble_block_matrix_kernel!( - field::F1, pattern::Patt, physics::Phys, ref_fe::R, - U::F2, U_old::F2, X::F3, state_old::S, state_new::S, props::P, t::T, dt::T, - conns::C, block_id::Int, func::Func + field::AbstractArray{<:Number, 1}, + conns::Connectivity, block_start_index::Int, block_el_level_size::Int, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, dt::T, + U::Solution, U_old::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - F1 <: AbstractArray{<:Number, 1}, - F2 <: AbstractField, - F3 <: AbstractField, - Func <: Function, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Patt <: SparseMatrixPattern, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } E = KA.@index(Global) @@ -133,7 +127,6 @@ KA.@kernel function _assemble_block_matrix_kernel!( for q in 1:num_quadrature_points(ref_fe) interps = _cell_interpolants(ref_fe, q) state_old_q = _quadrature_level_state(state_old, q, E) - # K_q, state_new_q = func(physics, interps, u_el, x_el, state_old_q, props_el, t, dt) K_q, state_new_q = func(physics, interps, x_el, t, dt, u_el, u_el_old, state_old_q, props_el) K_el = K_el + K_q for s in 1:length(state_old) @@ -141,15 +134,17 @@ KA.@kernel function _assemble_block_matrix_kernel!( end end - block_start_index = values(pattern.block_start_indices)[block_id] - block_el_level_size = values(pattern.block_el_level_sizes)[block_id] - start_id = block_start_index + - (E - 1) * block_el_level_size - end_id = start_id + block_el_level_size - 1 - ids = start_id:end_id - for (i, id) in enumerate(ids) - Atomix.@atomic field[id] += K_el.data[i] - end + # leaving here just in case + # block_start_index = values(pattern.block_start_indices)[block_id] + # block_el_level_size = values(pattern.block_el_level_sizes)[block_id] + # start_id = block_start_index + + # (E - 1) * block_el_level_size + # end_id = start_id + block_el_level_size - 1 + # ids = start_id:end_id + # for (i, id) in enumerate(ids) + # Atomix.@atomic field[id] += K_el.data[i] + # end + _assemble_element!(field, K_el, E, block_start_index, block_el_level_size) end # COV_EXCL_STOP @@ -161,28 +156,29 @@ using KernelAbstractions and Atomix for eliminating race conditions TODO add state variables and physics properties """ function _assemble_block_matrix!( - field::F1, pattern::Patt, physics::Phys, ref_fe::R, - U::F2, U_old::F2, X::F3, state_old::S, state_new::S, props::P, t::T, dt::T, - conns::C, block_id::Int, func::Func, backend::KA.Backend + backend::KA.Backend, + field::AbstractArray{<:Number, 1}, + conns::Connectivity, block_start_index::Int, block_el_level_size::Int, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, dt::T, + U::Solution, U_old::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - F1 <: AbstractArray{<:Number, 1}, - F2 <: AbstractField, - F3 <: AbstractField, - Func <: Function, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Patt <: SparseMatrixPattern, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } kernel! = _assemble_block_matrix_kernel!(backend) kernel!( - field, pattern, physics, ref_fe, - U, U_old, X, state_old, state_new, props, t, dt, - conns, block_id, func, ndrange=size(conns, 2) + field, + conns, block_start_index, block_el_level_size, + func, + physics, ref_fe, + X, t, dt, + U, U_old, + state_old, state_new, props, + ndrange=size(conns, 2) ) return nothing end diff --git a/src/assemblers/MatrixAction.jl b/src/assemblers/MatrixAction.jl index 4b779d2..c37545f 100644 --- a/src/assemblers/MatrixAction.jl +++ b/src/assemblers/MatrixAction.jl @@ -16,29 +16,33 @@ end $(TYPEDSIGNATURES) """ function assemble_matrix_action!( - # assembler, func, Uu, Vu, p storage, dof, func, Uu, Vu, p ) - # storage = assembler.stiffness_action_storage fill!(storage, zero(eltype(storage))) fspace = function_space(dof) t = current_time(p.times) Δt = time_step(p.times) _update_for_assembly!(p, dof, Uu, Vu) - for (b, (conns, block_physics, state_old, state_new, props)) in enumerate(zip( + for ( + conns, + block_physics, ref_fe, + state_old, state_new, props + ) in zip( values(fspace.elem_conns), - values(p.physics), + values(p.physics), values(fspace.ref_fes), values(p.state_old), values(p.state_new), values(p.properties) - )) - ref_fe = values(fspace.ref_fes)[b] - # backend = _check_backends(assembler, p.h1_field, p.h1_hvp_scratch_field, p.h1_coords, state_old, state_new, conns) + ) backend = KA.get_backend(p.h1_field) _assemble_block_matrix_action!( - storage, block_physics, ref_fe, - p.h1_field, p.h1_field_old, p.h1_hvp_scratch_field, p.h1_coords, state_old, state_new, props, t, Δt, - conns, b, func, - backend + backend, + storage, + conns, + func, + block_physics, ref_fe, + p.h1_coords, t, Δt, + p.h1_field, p.h1_field_old, p.h1_hvp_scratch_field, + state_old, state_new, props ) end end @@ -54,22 +58,18 @@ TODO improve typing of fields to ensure they mathc up in terms of function spaces """ function _assemble_block_matrix_action!( - field::F1, physics::Phys, ref_fe::R, - U::F2, U_old::F2, V::F3, X::F4, state_old::S, state_new::S, props::P, t::T, Δt::T, - conns::C, block_id::Int, func::Func, ::KA.CPU + ::KA.CPU, + field::AbstractField, + conns::Connectivity, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, Δt::T, + U::Solution, U_old::Solution, V::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - F1 <: AbstractField, - F2 <: AbstractField, - F3 <: AbstractField, - F4 <: AbstractField, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Func <: Function, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } for e in axes(conns, 2) x_el = _element_level_fields_flat(X, ref_fe, conns, e) @@ -89,7 +89,7 @@ function _assemble_block_matrix_action!( end end Kv_el = K_el * v_el - @views _assemble_element!(field, Kv_el, conns[:, e], e, block_id) + @views _assemble_element!(field, Kv_el, conns[:, e]) end end @@ -103,22 +103,17 @@ TODO mark const fields """ # COV_EXCL_START KA.@kernel function _assemble_block_matrix_action_kernel!( - field::F1, physics::Phys, ref_fe::R, - U::F2, U_old::F2, V::F3, X::F4, state_old::S, state_new::S, props::P, t::T, Δt::T, - conns::C, block_id::Int, func::Func + field::AbstractField, + conns::Connectivity, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, Δt::T, + U::Solution, U_old::Solution, V::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - F1 <: AbstractField, - F2 <: AbstractField, - F3 <: AbstractField, - F4 <: AbstractField, - Func <: Function, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } E = KA.@index(Global) @@ -160,28 +155,29 @@ using KernelAbstractions and Atomix for eliminating race conditions TODO add state variables and physics properties """ function _assemble_block_matrix_action!( - field::F1, physics::Phys, ref_fe::R, - U::F2, U_old::F2, V::F3, X::F4, state_old::S, state_new::S, props::P, t::T, Δt::T, - conns::C, block_id::Int, func::Func, backend::KA.Backend + backend::KA.Backend, + field::AbstractField, + conns::Connectivity, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, Δt::T, + U::Solution, U_old::Solution, V::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - F1 <: AbstractField, - F2 <: AbstractField, - F3 <: AbstractField, - F4 <: AbstractField, - Func <: Function, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } kernel! = _assemble_block_matrix_action_kernel!(backend) kernel!( - field, physics, ref_fe, - U, U_old, V, X, state_old, state_new, props, t, Δt, - conns, block_id, func, ndrange=size(conns, 2) + field, + conns, + func, + physics, ref_fe, + X, t, Δt, + U, U_old, V, + state_old, state_new, props, t, Δt, + ndrange=size(conns, 2) ) return nothing end diff --git a/src/assemblers/NeumannBC.jl b/src/assemblers/NeumannBC.jl index 77511a0..b813c09 100644 --- a/src/assemblers/NeumannBC.jl +++ b/src/assemblers/NeumannBC.jl @@ -17,20 +17,18 @@ end $(TYPEDSIGNATURES) """ function assemble_vector_neumann_bc!( - # assembler, Uu, p storage, dof, Uu, p ) - # storage = assembler.residual_storage # do not zero! - t = current_time(p.times) # TODO should below 2 methods calls be assumed to have # been conducted previously? _update_for_assembly!(p, dof, Uu) for bc in values(p.neumann_bcs.bc_caches) backend = KA.get_backend(bc) _assemble_block_vector_neumann_bc!( - storage, p.h1_field, p.h1_coords, t, - bc, backend + backend, storage, + p.h1_field, p.h1_coords, + bc ) end end @@ -39,15 +37,10 @@ end $(TYPEDSIGNATURES) """ function _assemble_block_vector_neumann_bc!( - field::F1, U::F2, X::F3, t::T, - bc::N, ::KA.CPU -) where { - F1 <: AbstractField, - F2 <: AbstractField, - F3 <: AbstractField, - T <: Number, - N <: NeumannBCContainer -} + ::KA.CPU, + field::AbstractField, U::AbstractField, X::AbstractField, + bc::NeumannBCContainer +) conns = bc.element_conns ref_fe = bc.ref_fe @@ -68,9 +61,8 @@ function _assemble_block_vector_neumann_bc!( R_el = R_el + JxW * Nvec * f_val end - block_id = 1 # doesn't matter for this method - el_id = bc.elements[e] - @views _assemble_element!(field, R_el, bc.surface_conns[:, e], el_id, block_id) + + @views _assemble_element!(field, R_el, bc.surface_conns[:, e]) end end @@ -82,15 +74,9 @@ TODO mark const fields """ # COV_EXCL_START KA.@kernel function _assemble_block_vector_neumann_bc_kernel!( - field::F1, U::F2, X::F3, t::T, - bc::N -) where { - F1 <: AbstractField, - F2 <: AbstractField, - F3 <: AbstractField, - T <: Number, - N <: NeumannBCContainer -} + field::AbstractField, U::AbstractField, X::AbstractField, + bc::NeumannBCContainer +) E = KA.@index(Global) conns = bc.element_conns @@ -130,18 +116,14 @@ end $(TYPEDSIGNATURES) """ function _assemble_block_vector_neumann_bc!( - field::F1, U::F2, X::F3, t::T, - bc::N, backend::KA.Backend -) where { - F1 <: AbstractField, - F2 <: AbstractField, - F3 <: AbstractField, - T <: Number, - N <: NeumannBCContainer -} + backend::KA.Backend, + field::AbstractField, + U::AbstractField, X::AbstractField, + bc::NeumannBCContainer +) kernel! = _assemble_block_vector_neumann_bc_kernel!(backend) kernel!( - field, U, X, t, bc, ndrange=size(bc.vals, 2) + field, U, X, bc, ndrange=size(bc.vals, 2) ) return nothing end diff --git a/src/assemblers/QuadratureQuantity.jl b/src/assemblers/QuadratureQuantity.jl index e1970a6..6bd0839 100644 --- a/src/assemblers/QuadratureQuantity.jl +++ b/src/assemblers/QuadratureQuantity.jl @@ -21,21 +21,28 @@ function assemble_quadrature_quantity!( t = current_time(p.times) Δt = time_step(p.times) _update_for_assembly!(p, dof, Uu) - for (b, (field, conns, block_physics, state_old, state_new, props)) in enumerate(zip( + for ( + block_storage, + conns, + block_physics, ref_fe, + state_old, state_new, props + ) in zip( values(storage), values(fspace.elem_conns), - values(p.physics), + values(p.physics), values(fspace.ref_fes), values(p.state_old), values(p.state_new), values(p.properties) - )) - ref_fe = values(fspace.ref_fes)[b] - # backend = _check_backends(assembler, p.h1_field, p.h1_field_old, p.h1_coords, state_old, state_new, conns) + ) backend = KA.get_backend(p.h1_field) _assemble_block_quadrature_quantity!( - field, block_physics, ref_fe, - p.h1_field, p.h1_field_old, p.h1_coords, state_old, state_new, props, t, Δt, - conns, b, func, - backend + backend, + block_storage, + conns, + func, + block_physics, ref_fe, + p.h1_coords, t, Δt, + p.h1_field, p.h1_field_old, + state_old, state_new, props ) end end @@ -51,22 +58,18 @@ TODO improve typing of fields to ensure they mathc up in terms of function spaces """ function _assemble_block_quadrature_quantity!( - field::F1, physics::Phys, ref_fe::R, - U::F2, U_old::F2, X::F3, state_old::S, state_new::S, props::P, t::T, Δt::T, - conns::C, block_id::Int, func::Func, ::KA.CPU + ::KA.CPU, + field::AbstractArray, + conns::Connectivity, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, Δt::T, + U::Solution, U_old::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - # F1 <: AbstractArray{<:Number, 2}, - F1 <: AbstractMatrix, - F2 <: AbstractField, - F3 <: AbstractField, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Func <: Function, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } for e in axes(conns, 2) x_el = _element_level_fields_flat(X, ref_fe, conns, e) @@ -95,22 +98,17 @@ TODO mark const fields """ # COV_EXCL_START KA.@kernel function _assemble_block_quadrature_quantity_kernel!( - field::F1, physics::Phys, ref_fe::R, - U::F2, U_old::F2, X::F3, state_old::S, state_new::S, props::P, t::T, Δt::T, - conns::C, block_id::Int, func::Func + field::AbstractArray, + conns::Connectivity, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, Δt::T, + U::Solution, U_old::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - # F1 <: AbstractArray{<:Number, 2}, - F1 <: AbstractMatrix, - F2 <: AbstractField, - F3 <: AbstractField, - Func <: Function, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } # Q, E = KA.@index(Global, NTuple) E = KA.@index(Global) @@ -139,28 +137,29 @@ using KernelAbstractions and Atomix for eliminating race conditions TODO add state variables and physics properties """ function _assemble_block_quadrature_quantity!( - field::F1, physics::Phys, ref_fe::R, - U::F2, U_old::F2, X::F3, state_old::S, state_new::S, props::P, t::T, Δt::T, - conns::C, block_id::Int, func::Func, backend::KA.Backend + backend::KA.Backend, + field::AbstractArray, + conns::Connectivity, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, Δt::T, + U::Solution, U_old::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - # F1 <: AbstractArray{<:Number, 2}, - F1 <: AbstractMatrix, - F2 <: AbstractField, - F3 <: AbstractField, - Func <: Function, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } kernel! = _assemble_block_quadrature_quantity_kernel!(backend) kernel!( - field, physics, ref_fe, - U, U_old, X, state_old, state_new, props, t, Δt, - conns, block_id, func, ndrange=size(field, 2) + field, + conns, + func, + physics, ref_fe, + X, t, Δt, + U, U_old, + state_old, state_new, props, + ndrange=size(field, 2) ) # KA.synchronize(backend) return nothing diff --git a/src/assemblers/SparseMatrixAssembler.jl b/src/assemblers/SparseMatrixAssembler.jl index deba0fc..e58ecfb 100644 --- a/src/assemblers/SparseMatrixAssembler.jl +++ b/src/assemblers/SparseMatrixAssembler.jl @@ -8,7 +8,7 @@ struct SparseMatrixAssembler{ Condensed, NumArrDims, NumFields, - BlockPattern <: NamedTuple, + # BlockPattern <: NamedTuple, IV <: AbstractArray{Int, 1}, RV <: AbstractArray{Float64, 1}, Var <: AbstractFunction, @@ -16,7 +16,8 @@ struct SparseMatrixAssembler{ QuadratureStorage <: NamedTuple } <: AbstractAssembler{DofManager{Condensed, Int, IV, Var}} dof::DofManager{Condensed, Int, IV, Var} - matrix_pattern::SparseMatrixPattern{IV, BlockPattern, RV} + # matrix_pattern::SparseMatrixPattern{IV, BlockPattern, RV} + matrix_pattern::SparseMatrixPattern{IV, RV} vector_pattern::SparseVectorPattern{IV} constraint_storage::RV damping_storage::RV @@ -109,26 +110,6 @@ function Base.show(io::IO, asm::SparseMatrixAssembler) println(io, " ", asm.dof) end -""" -$(TYPEDSIGNATURES) -Specialization of of ```_assemble_element!``` for ```SparseMatrixAssembler```. -""" -function _assemble_element!( - pattern::SparseMatrixPattern, storage, K_el::SMatrix, el_id::Int, block_id::Int -) - # figure out ids needed to update - block_start_index = values(pattern.block_start_indices)[block_id] - block_el_level_size = values(pattern.block_el_level_sizes)[block_id] - start_id = block_start_index + - (el_id - 1) * block_el_level_size - end_id = start_id + block_el_level_size - 1 - ids = start_id:end_id - - # get appropriate storage and update values - @views storage[ids] += K_el[:] - return nothing -end - # TODO this only work on CPU right now function _adjust_matrix_entries_for_constraints!( A::SparseMatrixCSC, constraint_storage, ::KA.CPU; diff --git a/src/assemblers/SparsityPatterns.jl b/src/assemblers/SparsityPatterns.jl index 4d364a8..8fdcf4a 100644 --- a/src/assemblers/SparsityPatterns.jl +++ b/src/assemblers/SparsityPatterns.jl @@ -7,14 +7,16 @@ case where you want to eliminate fixed-dofs or not. """ struct SparseMatrixPattern{ I <: AbstractArray{Int, 1}, - B, + # B, R <: AbstractArray{Float64, 1} } Is::I Js::I unknown_dofs::I - block_start_indices::B - block_el_level_sizes::B + # block_start_indices::B + # block_el_level_sizes::B + block_start_indices::Vector{Int} + block_el_level_sizes::Vector{Int} # cache arrays klasttouch::I csrrowptr::I @@ -60,9 +62,9 @@ function SparseMatrixPattern(dof::DofManager) end # convert to NamedTuples so it's easy to index - block_syms = keys(fspace.ref_fes) - block_start_indices = NamedTuple{block_syms}(tuple(block_start_indices)...) - block_el_level_sizes = NamedTuple{block_syms}(tuple(block_el_level_sizes)...) + # block_syms = keys(fspace.ref_fes) + # block_start_indices = NamedTuple{block_syms}(tuple(block_start_indices)...) + # block_el_level_sizes = NamedTuple{block_syms}(tuple(block_el_level_sizes)...) # setup pre-allocated arrays based on number of entries found above Is = Vector{Int64}(undef, n_entries) @@ -112,8 +114,8 @@ function Adapt.adapt_structure(to, asm::SparseMatrixPattern) Is = adapt(to, asm.Is) Js = adapt(to, asm.Js) unknown_dofs = adapt(to, asm.unknown_dofs) - block_start_indices = adapt(to, asm.block_start_indices) - block_el_level_sizes = adapt(to, asm.block_el_level_sizes) + # block_start_indices = adapt(to, asm.block_start_indices) + # block_el_level_sizes = adapt(to, asm.block_el_level_sizes) # klasttouch = adapt(to, asm.klasttouch) csrrowptr = adapt(to, asm.csrrowptr) @@ -125,7 +127,7 @@ function Adapt.adapt_structure(to, asm::SparseMatrixPattern) cscnzval = adapt(to, asm.cscnzval) return SparseMatrixPattern( Is, Js, - unknown_dofs, block_start_indices, block_el_level_sizes, + unknown_dofs, asm.block_start_indices, asm.block_el_level_sizes, klasttouch, csrrowptr, csrcolval, csrnzval, csccolptr, cscrowval, cscnzval ) @@ -140,6 +142,29 @@ function SparseArrays.sparse!(pattern::SparseMatrixPattern, storage) ) end +""" +$(TYPEDSIGNATURES) +Specialization of of ```_assemble_element!``` for ```SparseMatrixAssembler```. +""" +function _assemble_element!( + storage, K_el::SMatrix, + el_id::Int, + block_start_index::Int, block_el_level_size::Int +) + # figure out ids needed to update + start_id = block_start_index + + (el_id - 1) * block_el_level_size + end_id = start_id + block_el_level_size - 1 + ids = start_id:end_id + + # get appropriate storage and update values + # @views storage[ids] += K_el[:] + for (i, id) in enumerate(ids) + storage[id] += K_el.data[i] + end + return nothing +end + num_entries(s::SparseMatrixPattern) = length(s.Is) # NOTE this methods assumes that dof is up to date diff --git a/src/assemblers/Vector.jl b/src/assemblers/Vector.jl index aede4f0..caf0c92 100644 --- a/src/assemblers/Vector.jl +++ b/src/assemblers/Vector.jl @@ -22,20 +22,26 @@ function assemble_vector!( t = current_time(p.times) Δt = time_step(p.times) _update_for_assembly!(p, dof, Uu) - for (b, (conns, block_physics, state_old, state_new, props)) in enumerate(zip( + for ( + conns, + block_physics, ref_fe, + state_old, state_new, props + ) in zip( values(fspace.elem_conns), - values(p.physics), + values(p.physics), values(fspace.ref_fes), values(p.state_old), values(p.state_new), values(p.properties) - )) - ref_fe = values(fspace.ref_fes)[b] + ) backend = KA.get_backend(p.h1_field) _assemble_block_vector!( - storage, block_physics, ref_fe, - p.h1_field, p.h1_field_old, p.h1_coords, state_old, state_new, props, t, Δt, - conns, b, + backend, + storage, + conns, func, - backend + block_physics, ref_fe, + p.h1_coords, t, Δt, + p.h1_field, p.h1_field_old, + state_old, state_new, props ) end @@ -53,22 +59,18 @@ TODO improve typing of fields to ensure they mathc up in terms of function spaces """ function _assemble_block_vector!( - field::F1, physics::Phys, ref_fe::R, - U::F2, U_old::F2, X::F3, state_old::S, state_new::S, props::P, t::T, Δt::T, - conns::C, block_id::Int, - func::Func, ::KA.CPU + ::KA.CPU, + field::AbstractField, + conns::Connectivity, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, Δt::T, + U::Solution, U_old::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - F1 <: AbstractField, - F2 <: AbstractField, - F3 <: AbstractField, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Func <: Function, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } for e in axes(conns, 2) x_el = _element_level_fields_flat(X, ref_fe, conns, e) @@ -80,7 +82,6 @@ function _assemble_block_vector!( for q in 1:num_quadrature_points(ref_fe) interps = _cell_interpolants(ref_fe, q) state_old_q = _quadrature_level_state(state_old, q, e) - # R_q, state_new_q = func(physics, interps, u_el, x_el, state_old_q, props_el, t, Δt) R_q, state_new_q = func(physics, interps, x_el, t, Δt, u_el, u_el_old, state_old_q, props_el) R_el = R_el + R_q # update state here @@ -89,7 +90,7 @@ function _assemble_block_vector!( end end - @views _assemble_element!(field, R_el, conns[:, e], e, block_id) + @views _assemble_element!(field, R_el, conns[:, e]) end end @@ -103,22 +104,17 @@ TODO mark const fields """ # COV_EXCL_START KA.@kernel function _assemble_block_vector_kernel!( - field::F1, physics::Phys, ref_fe::R, - U::F2, U_old::F2, X::F3, state_old::S, state_new::S, props::P, t::T, Δt::T, - conns::C, block_id::Int, - func::Func + field::AbstractField, + conns::Connectivity, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, Δt::T, + U::Solution, U_old::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - F1 <: AbstractField, - F2 <: AbstractField, - F3 <: AbstractField, - Func <: Function, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } E = KA.@index(Global) @@ -131,7 +127,6 @@ KA.@kernel function _assemble_block_vector_kernel!( for q in 1:num_quadrature_points(ref_fe) interps = _cell_interpolants(ref_fe, q) state_old_q = _quadrature_level_state(state_old, q, E) - # R_q, state_new_q = func(physics, interps, u_el, x_el, state_old_q, props_el, t, Δt) R_q, state_new_q = func(physics, interps, x_el, t, Δt, u_el, u_el_old, state_old_q, props_el) R_el = R_el + R_q # update state here @@ -160,29 +155,29 @@ using KernelAbstractions and Atomix for eliminating race conditions TODO add state variables and physics properties """ function _assemble_block_vector!( - field::F1, physics::Phys, ref_fe::R, - U::F2, U_old::F2, X::F3, state_old::S, state_new::S, props::P, t::T, Δt::T, - conns::C, block_id::Int, - func::Func, backend::KA.Backend + backend::KA.Backend, + field::AbstractField, + conns::Connectivity, + func::Function, + physics::AbstractPhysics, ref_fe::ReferenceFE, + X::AbstractField, t::T, Δt::T, + U::Solution, U_old::Solution, + state_old::S, state_new::S, props::AbstractArray ) where { - C <: Connectivity, - F1 <: AbstractField, - F2 <: AbstractField, - F3 <: AbstractField, - Func <: Function, - # P <: Union{<:SVector, <:L2ElementField}, - P <: AbstractArray, - Phys <: AbstractPhysics, - R <: ReferenceFE, - S <: L2QuadratureField, - T <: Number + T <: Number, + Solution <: AbstractField, + S <: L2QuadratureField } kernel! = _assemble_block_vector_kernel!(backend) kernel!( - field, physics, ref_fe, - U, U_old, X, state_old, state_new, props, t, Δt, - conns, block_id, - func, ndrange=size(conns, 2) + field, + conns, + func, + physics, ref_fe, + X, t, Δt, + U, U_old, + state_old, state_new, props, + ndrange=size(conns, 2) ) return nothing end