Skip to content

Commit

Permalink
First round of fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
JordiManyer committed Aug 13, 2023
1 parent 1a1c47a commit 9b1620c
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 66 deletions.
22 changes: 16 additions & 6 deletions src/FESpaces.jl
Original file line number Diff line number Diff line change
Expand Up @@ -600,27 +600,37 @@ FESpaces.get_vector_builder(a::DistributedSparseMatrixAssembler) = a.vector_buil
FESpaces.get_assembly_strategy(a::DistributedSparseMatrixAssembler) = a.strategy

function FESpaces.symbolic_loop_matrix!(A,a::DistributedSparseMatrixAssembler,matdata)
map(symbolic_loop_matrix!,local_views(A,a.test_dofs_gids_prange,a.trial_dofs_gids_prange),a.assems,matdata)
rows = get_rows(a)
cols = get_cols(a)
map(symbolic_loop_matrix!,local_views(A,rows,cols),local_views(a),matdata)
end

function FESpaces.numeric_loop_matrix!(A,a::DistributedSparseMatrixAssembler,matdata)
map(numeric_loop_matrix!,local_views(A,a.test_dofs_gids_prange,a.trial_dofs_gids_prange),a.assems,matdata)
rows = get_rows(a)
cols = get_cols(a)
map(numeric_loop_matrix!,local_views(A,rows,cols),local_views(a),matdata)
end

function FESpaces.symbolic_loop_vector!(b,a::DistributedSparseMatrixAssembler,vecdata)
map(symbolic_loop_vector!,local_views(b,a.test_dofs_gids_prange),a.assems,vecdata)
rows = get_rows(a)
map(symbolic_loop_vector!,local_views(b,rows),local_views(a),vecdata)
end

function FESpaces.numeric_loop_vector!(b,a::DistributedSparseMatrixAssembler,vecdata)
map(numeric_loop_vector!,local_views(b,a.test_dofs_gids_prange),a.assems,vecdata)
rows = get_rows(a)
map(numeric_loop_vector!,local_views(b,rows),local_views(a),vecdata)
end

function FESpaces.symbolic_loop_matrix_and_vector!(A,b,a::DistributedSparseMatrixAssembler,data)
map(symbolic_loop_matrix_and_vector!,local_views(A,a.test_dofs_gids_prange,a.trial_dofs_gids_prange),local_views(b,a.test_dofs_gids_prange),a.assems,data)
rows = get_rows(a)
cols = get_cols(a)
map(symbolic_loop_matrix_and_vector!,local_views(A,rows,cols),local_views(b,rows),local_views(a),data)
end

function FESpaces.numeric_loop_matrix_and_vector!(A,b,a::DistributedSparseMatrixAssembler,data)
map(numeric_loop_matrix_and_vector!,local_views(A,a.test_dofs_gids_prange,a.trial_dofs_gids_prange),local_views(b,a.test_dofs_gids_prange),a.assems,data)
rows = get_rows(a)
cols = get_cols(a)
map(numeric_loop_matrix_and_vector!,local_views(A,rows,cols),local_views(b,rows),local_views(a),data)
end

# Parallel Assembly strategies
Expand Down
89 changes: 29 additions & 60 deletions src/MultiField.jl
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,9 @@ end

# BlockSparseMatrixAssemblers

const DistributedBlockSparseMatrixAssembler{NB,NV,SB,P} =
MultiField.BlockSparseMatrixAssembler{NB,NV,SB,P,<:DistributedSparseMatrixAssembler}

function FESpaces.SparseMatrixAssembler(
local_mat_type,
local_vec_type,
Expand Down Expand Up @@ -433,7 +436,7 @@ end
function local_views(a::MultiField.BlockSparseMatrixAssembler{NB,NV,SB,P}) where {NB,NV,SB,P}
assems = a.block_assemblers
parts = get_part_ids(local_views(first(assems)))
map_parts(parts) do p
map(parts) do p
idx = CartesianIndices(axes(assems))
block_assems = map(idx) do I
get_part(local_views(assems[I]),p)
Expand All @@ -444,7 +447,7 @@ end

function local_views(a::MatrixBlock,rows,cols)
parts = get_part_ids(local_views(first(a.array)))
map_parts(parts) do p
map(parts) do p
idx = CartesianIndices(axes(a))
array = map(idx) do I
get_part(local_views(a[I],rows[I[1]],cols[I[2]]),p)
Expand All @@ -455,7 +458,7 @@ end

function local_views(a::VectorBlock,rows)
parts = get_part_ids(local_views(first(a.array)))
map_parts(parts) do p
map(parts) do p
idx = CartesianIndices(axes(a))
array = map(idx) do I
get_part(local_views(a[I],rows[I]),p)
Expand All @@ -464,48 +467,46 @@ function local_views(a::VectorBlock,rows)
end
end

function local_views(a::MultiField.ArrayBlockView,axes...)
function local_views(a::ArrayBlockView,axes...)
array = local_views(a.array,axes...)
map_parts(array) do array
map(array) do array
MultiField.ArrayBlockView(array,a.block_map)
end
end


#! The following functions could be avoided if we created an abstract superclass for
#! DistributedSparseMatrixAssembler
function FESpaces.symbolic_loop_matrix!(A,a::MultiField.BlockSparseMatrixAssembler,matdata::AbstractPData)
#! The following could be avoided if DistributedBlockSparseMatrixAssembler <: DistributedSparseMatrixAssembler
function FESpaces.symbolic_loop_matrix!(A,a::DistributedBlockSparseMatrixAssembler,matdata)
rows = get_rows(a)
cols = get_cols(a)
map_parts(symbolic_loop_matrix!,local_views(A,rows,cols),local_views(a),matdata)
map(symbolic_loop_matrix!,local_views(A,rows,cols),local_views(a),matdata)
end

function FESpaces.numeric_loop_matrix!(A,a::MultiField.BlockSparseMatrixAssembler,matdata::AbstractPData)
function FESpaces.numeric_loop_matrix!(A,a::DistributedBlockSparseMatrixAssembler,matdata)
rows = get_rows(a)
cols = get_cols(a)
map_parts(numeric_loop_matrix!,local_views(A,rows,cols),local_views(a),matdata)
map(numeric_loop_matrix!,local_views(A,rows,cols),local_views(a),matdata)
end

function FESpaces.symbolic_loop_vector!(b,a::MultiField.BlockSparseMatrixAssembler,vecdata::AbstractPData)
function FESpaces.symbolic_loop_vector!(b,a::DistributedBlockSparseMatrixAssembler,vecdata)
rows = get_rows(a)
map_parts(symbolic_loop_vector!,local_views(b,rows),local_views(a),vecdata)
map(symbolic_loop_vector!,local_views(b,rows),local_views(a),vecdata)
end

function FESpaces.numeric_loop_vector!(b,a::MultiField.BlockSparseMatrixAssembler,vecdata::AbstractPData)
function FESpaces.numeric_loop_vector!(b,a::DistributedBlockSparseMatrixAssembler,vecdata)
rows = get_rows(a)
map_parts(numeric_loop_vector!,local_views(b,rows),local_views(a),vecdata)
map(numeric_loop_vector!,local_views(b,rows),local_views(a),vecdata)
end

function FESpaces.symbolic_loop_matrix_and_vector!(A,b,a::MultiField.BlockSparseMatrixAssembler,data::AbstractPData)
function FESpaces.symbolic_loop_matrix_and_vector!(A,b,a::DistributedBlockSparseMatrixAssembler,data)
rows = get_rows(a)
cols = get_cols(a)
map_parts(symbolic_loop_matrix_and_vector!,local_views(A,rows,cols),local_views(b,rows),local_views(a),data)
map(symbolic_loop_matrix_and_vector!,local_views(A,rows,cols),local_views(b,rows),local_views(a),data)
end

function FESpaces.numeric_loop_matrix_and_vector!(A,b,a::MultiField.BlockSparseMatrixAssembler,data::AbstractPData)
function FESpaces.numeric_loop_matrix_and_vector!(A,b,a::DistributedBlockSparseMatrixAssembler,data)
rows = get_rows(a)
cols = get_cols(a)
map_parts(numeric_loop_matrix_and_vector!,local_views(A,rows,cols),local_views(b,rows),local_views(a),data)
map(numeric_loop_matrix_and_vector!,local_views(A,rows,cols),local_views(b,rows),local_views(a),data)
end

#! The following is horrible (see dicussion in PR) but necessary for the moment. We will be
Expand All @@ -518,21 +519,21 @@ function Algebra.create_from_nz(a::ArrayBlock{<:DistributedAllocationCOO{<:Fully
end

function _fa_create_from_nz_temporary_fix(a::DistributedAllocationCOO{<:FullyAssembledRows})
parts = get_part_ids(local_views(a))
parts = get_part_ids(local_views(a))

rdofs = a.rows # dof ids of the test space
cdofs = a.cols # dof ids of the trial space
ngrdofs = length(rdofs)
ngcdofs = length(cdofs)
nordofs = map_parts(num_oids,rdofs.partition)
nocdofs = map_parts(num_oids,cdofs.partition)
first_grdof = map_parts(first_gdof_from_ids,rdofs.partition)
first_gcdof = map_parts(first_gdof_from_ids,cdofs.partition)
nordofs = map(num_oids,rdofs.partition)
nocdofs = map(num_oids,cdofs.partition)
first_grdof = map(first_gdof_from_ids,rdofs.partition)
first_gcdof = map(first_gdof_from_ids,cdofs.partition)
cneigs_snd = cdofs.exchanger.parts_snd
cneigs_rcv = cdofs.exchanger.parts_rcv

hcol_to_gid = map_parts(part -> part.lid_to_gid[part.hid_to_lid], cdofs.partition)
hcol_to_part = map_parts(part -> part.lid_to_part[part.hid_to_lid], cdofs.partition)
hcol_to_gid = map(part -> part.lid_to_gid[part.hid_to_lid], cdofs.partition)
hcol_to_part = map(part -> part.lid_to_part[part.hid_to_lid], cdofs.partition)

rows = PRange(
parts,
Expand Down Expand Up @@ -560,39 +561,7 @@ function _fa_create_from_nz_temporary_fix(a::DistributedAllocationCOO{<:FullyAss

b = change_axes(a,(rows,cols))

values = map_parts(Algebra.create_from_nz,local_views(b))
values = map(Algebra.create_from_nz,local_views(b))
exchanger = empty_exchanger(parts)
return PSparseMatrix(values,rows,cols,exchanger)
end

"""
function Algebra.nz_allocation(a::ArrayBlock{<:DistributedCounterCOO})
array = map(Algebra.nz_allocation,a.array)
match_block_indexes!(array)
return ArrayBlock(array,a.touched)
end
function match_block_indexes!(allocators::Vector{<:DistributedAllocationCOO})
return allocators
end
function match_block_indexes!(allocators::Matrix{<:DistributedAllocationCOO})
s = size(allocators)
# Get an AbstractPData containing in each part the the matrix of local allocators
parts = get_part_ids(local_views(first(allocators)))
allocs = map_parts(parts) do p
idx = CartesianIndices(s)
allocs = map(idx) do I
get_part(local_views(a[I],rows[I[1]],cols[I[2]]),p)
end
return allocs
end
# Accumulate the index sets for each
map_parts()
for block_row in 1:s[1]
end
end
"""

0 comments on commit 9b1620c

Please sign in to comment.