Merge pull request #143 from gridap/fix_issue_142

Fix issue 142

Author: amartinhuertas

NEWS.md

@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Added redistribution for MultiFieldFESpaces. Since PR [#140](https://github.com/gridap/GridapDistributed.jl/pull/140).
 
+### Fixed 
+
+- Fixed issue [#142](https://github.com/gridap/GridapDistributed.jl/issues/142). Since PR [#142](https://github.com/gridap/GridapDistributed.jl/pull/142).
+
 ## [0.3.5] - 2023-12-04
 
 ### Added

Project.toml

@@ -17,7 +17,7 @@ WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 [compat]
 BlockArrays = "0.16.38"
 FillArrays = "0.8.4,1"
-Gridap = "0.17.21"
+Gridap = "0.17.23"
 LinearAlgebra = "1.3"
 MPI = "0.16, 0.17, 0.18, 0.19, 0.20"
 PartitionedArrays = "0.3.3"

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,18 +704,18 @@ 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)
     # now we can assemble contributions
     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,32 +762,39 @@ 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
 
 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,20 +820,21 @@ 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)
   end
   return b
 end
 
-
 # Common Assembly Utilities
-
 function first_gdof_from_ids(ids)
   if own_length(ids) == 0
     return 1

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)

test/issue_142.jl

@@ -0,0 +1,35 @@
+using Gridap
+using PartitionedArrays
+using GridapDistributed
+using Test
+
+function main(distribute,rank_partition)
+  DX             = 1000.0
+  DY             = 1000.0
+  order          = 0
+  n_els_x        = 4
+  n_els_y        = 4
+  dx             = DX/n_els_x
+  domain         = (0,DX,0,DY)
+  cell_partition = (n_els_x,n_els_y)
+  ranks          = distribute(LinearIndices((prod(rank_partition),)))
+
+  model = CartesianDiscreteModel(ranks,rank_partition,domain,cell_partition; isperiodic=(false,false))
+  Ω   = Triangulation(model)
+  dΩ  = Measure(Ω, 5*(order+1))
+  Γ   = SkeletonTriangulation(model)
+  dΓ  = Measure(Γ, 5*(order+1))
+
+  Q = FESpace(model, ReferenceFE(lagrangian, Float64, order), conformity=:L2)
+  P = TrialFESpace(Q)
+
+  # initial conditions
+  ph=FEFunction(P,prand(partition(Q.gids)))
+
+  b(q)   = ∫(jump(ph)*mean(q))dΓ
+  m(p,q) = ∫(p*q)dΩ
+  op = AffineFEOperator(m, b, P, Q)
+  b=assemble_vector(b(get_fe_basis(Q)),P)
+  tol=1.0e-12
+  @test norm(op.op.vector-b)/norm(b) < tol
+end 

test/sequential/issue_142.jl

@@ -0,0 +1,4 @@
+include("../issue_142.jl")
+with_debug() do distribute
+    main(distribute,(2,1))
+end

test/sequential/runtests.jl

@@ -12,6 +12,8 @@ using Test
 
 @time @testset "MultiField" begin include("MultiFieldTests.jl") end
 
+@time @testset "issue_142" begin include("issue_142.jl") end
+
 @time @testset "Poisson" begin include("PoissonTests.jl") end
 
 @time @testset "PLaplacian" begin include("PLaplacianTests.jl") end