diff --git a/src/Algebra.jl b/src/Algebra.jl index 45244ec..e211bb8 100644 --- a/src/Algebra.jl +++ b/src/Algebra.jl @@ -507,17 +507,17 @@ end function Algebra.create_from_nz(a::DistributedAllocationCOO{<:SubAssembledRows}) f(x) = nothing - A, = _sa_create_from_nz_with_callback(f,f,a) + A, = _sa_create_from_nz_with_callback(f,f,a,nothing) return A end function Algebra.create_from_nz(a::ArrayBlock{<:DistributedAllocationCOO{<:SubAssembledRows}}) f(x) = nothing - A, = _sa_create_from_nz_with_callback(f,f,a) + A, = _sa_create_from_nz_with_callback(f,f,a,nothing) return A end -function _sa_create_from_nz_with_callback(callback,async_callback,a) +function _sa_create_from_nz_with_callback(callback,async_callback,a,b) # Recover some data I,J,V = get_allocations(a) test_dofs_gids_prange = get_test_gids(a) @@ -538,7 +538,10 @@ function _sa_create_from_nz_with_callback(callback,async_callback,a) # Here we can overlap computations # This is a good place to overlap since # sending the matrix rows is a lot of data - b = callback(rows) + if !isa(b,Nothing) + bprange=_setup_prange_from_pvector_allocation(b) + b = callback(bprange) + end # Wait the transfer to finish wait(t) @@ -701,10 +704,10 @@ end function Algebra.create_from_nz( a::DistributedAllocationCOO{<:SubAssembledRows}, - c_fespace::PVectorAllocationTrackOnlyValues{<:SubAssembledRows}) + b::PVectorAllocationTrackTouchedAndValues) function callback(rows) - _rhs_callback(c_fespace,rows) + _rhs_callback(b,rows) end function async_callback(b) @@ -712,7 +715,7 @@ function Algebra.create_from_nz( assemble!(b) end - A,b = _sa_create_from_nz_with_callback(callback,async_callback,a) + A,b = _sa_create_from_nz_with_callback(callback,async_callback,a,b) return A,b end @@ -759,24 +762,31 @@ end nothing end + +function _setup_touched_and_allocations_arrays(values) + touched = map(values) do values + fill!(Vector{Bool}(undef,length(values)),false) + end + allocations = map(values,touched) do values,touched + ArrayAllocationTrackTouchedAndValues(touched,values) + end + touched, allocations +end + function Arrays.nz_allocation(a::DistributedCounterCOO{<:SubAssembledRows}, b::PVectorCounter{<:SubAssembledRows}) A = nz_allocation(a) dofs = b.test_dofs_gids_prange values = map(nz_allocation,b.counters) - B = PVectorAllocationTrackOnlyValues(b.par_strategy,values,dofs) + touched,allocations=_setup_touched_and_allocations_arrays(values) + B = PVectorAllocationTrackTouchedAndValues(allocations,values,dofs) return A,B end function Arrays.nz_allocation(a::PVectorCounter{<:SubAssembledRows}) dofs = a.test_dofs_gids_prange values = map(nz_allocation,a.counters) - touched = map(values) do values - fill!(Vector{Bool}(undef,length(values)),false) - end - allocations = map(values,touched) do values,touched - ArrayAllocationTrackTouchedAndValues(touched,values) - end + touched,allocations=_setup_touched_and_allocations_arrays(values) return PVectorAllocationTrackTouchedAndValues(allocations,values,dofs) end @@ -784,7 +794,7 @@ function local_views(a::PVectorAllocationTrackTouchedAndValues) a.allocations end -function Algebra.create_from_nz(a::PVectorAllocationTrackTouchedAndValues) +function _setup_prange_from_pvector_allocation(a::PVectorAllocationTrackTouchedAndValues) test_dofs_prange = a.test_dofs_gids_prange # dof ids of the test space ngrdofs = length(test_dofs_prange) @@ -810,10 +820,13 @@ function Algebra.create_from_nz(a::PVectorAllocationTrackTouchedAndValues) gids_ghost_to_global, gids_ghost_to_owner = map( find_gid_and_owner,I_ghost_lids_to_dofs_ghost_lids,indices) |> tuple_of_arrays - rows = _setup_prange_impl_(ngrdofs,indices,gids_ghost_to_global,gids_ghost_to_owner) + _setup_prange_impl_(ngrdofs,indices,gids_ghost_to_global,gids_ghost_to_owner) +end + +function Algebra.create_from_nz(a::PVectorAllocationTrackTouchedAndValues) + rows = _setup_prange_from_pvector_allocation(a) b = _rhs_callback(a,rows) t2 = assemble!(b) - # Wait the transfer to finish if t2 !== nothing wait(t2) @@ -821,9 +834,7 @@ function Algebra.create_from_nz(a::PVectorAllocationTrackTouchedAndValues) return b end - # Common Assembly Utilities - function first_gdof_from_ids(ids) if own_length(ids) == 0 return 1 diff --git a/src/FESpaces.jl b/src/FESpaces.jl index cd922dd..8972ff4 100644 --- a/src/FESpaces.jl +++ b/src/FESpaces.jl @@ -621,7 +621,10 @@ end function FESpaces.symbolic_loop_matrix_and_vector!(A,b,a::DistributedSparseMatrixAssembler,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) + Aviews=local_views(A,rows,cols) + bviews=local_views(b,rows) + aviews=local_views(a) + map(symbolic_loop_matrix_and_vector!,Aviews,bviews,aviews,data) end function FESpaces.numeric_loop_matrix_and_vector!(A,b,a::DistributedSparseMatrixAssembler,data)