Skip to content

cleaning up assemblers typing among other things. #207

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cmhamel merged 1 commit into from
Nov 28, 2025
Merged

cleaning up assemblers typing among other things. #207

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 1 addition & 28 deletions src/assemblers/Assemblers.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
"""
Expand Down
144 changes: 70 additions & 74 deletions src/assemblers/Matrix.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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)

Expand All @@ -133,23 +127,24 @@ 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)
state_new[s, q, E] = state_new_q[s]
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

Expand All @@ -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
Loading
Loading