Merge pull request #101 from fverdugo/performance

Performance

Author: fverdugo

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.3.2] - 2023-05-10
+
+### Fixed
+
+- Performance improvements.
+
 ## [0.3.1] - 2023-03-17
 
 ### Fixed

Project.toml

@@ -1,7 +1,7 @@
 name = "PartitionedArrays"
 uuid = "5a9dfac6-5c52-46f7-8278-5e2210713be9"
 authors = ["Francesc Verdugo <[email protected]> and contributors"]
-version = "0.3.1"
+version = "0.3.2"
 
 [deps]
 CircularArrays = "7a955b69-7140-5f4e-a0ed-f168c5e2e749"
@@ -18,7 +18,7 @@ SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1"
 [compat]
 CircularArrays = "1"
 Distances = "0.10"
-FillArrays = "0.10, 0.11, 0.12, 0.13"
+FillArrays = "0.10, 0.11, 0.12, 0.13, 1"
 IterativeSolvers = "0.9"
 MPI = "0.16, 0.17, 0.18, 0.19, 0.20"
 SparseMatricesCSR = "0.6"

src/mpi_array.jl

@@ -196,18 +196,55 @@ end
 function Base.map(f,args::Vararg{MPIArray,N}) where N
     a = first(args)
     @assert all(i->size(a)==size(i),args)
-    item = f(map(i->i.item,args)...)
+    item = f(map(i->i.item_ref[],args)...)
+    MPIArray(item,a.comm,a.size)
+end
+
+function Base.map(f,a::MPIArray)
+    item = f(a.item_ref[])
     MPIArray(item,a.comm,a.size)
 end
 
 function Base.map(f,args::Vararg{MPIArray,2})
     a = first(args)
     @assert all(i->size(a)==size(i),args)
-    t1,t2 = map(i->i.item,args)
+    t1,t2 = map(i->i.item_ref[],args)
     item = f(t1,t2)
     MPIArray(item,a.comm,a.size)
 end
 
+function Base.map(f,args::Vararg{MPIArray,3})
+    a = first(args)
+    @assert all(i->size(a)==size(i),args)
+    t1,t2,t3 = map(i->i.item_ref[],args)
+    item = f(t1,t2,t3)
+    MPIArray(item,a.comm,a.size)
+end
+
+function Base.map(f,args::Vararg{MPIArray,4})
+    a = first(args)
+    @assert all(i->size(a)==size(i),args)
+    t1,t2,t3,t4 = map(i->i.item_ref[],args)
+    item = f(t1,t2,t3,t4)
+    MPIArray(item,a.comm,a.size)
+end
+
+function Base.map(f,args::Vararg{MPIArray,5})
+    a = first(args)
+    @assert all(i->size(a)==size(i),args)
+    t1,t2,t3,t4,t5 = map(i->i.item_ref[],args)
+    item = f(t1,t2,t3,t4,t5)
+    MPIArray(item,a.comm,a.size)
+end
+
+function Base.map(f,args::Vararg{MPIArray,6})
+    a = first(args)
+    @assert all(i->size(a)==size(i),args)
+    t1,t2,t3,t4,t5,t6 = map(i->i.item_ref[],args)
+    item = f(t1,t2,t3,t4,t5,t6)
+    MPIArray(item,a.comm,a.size)
+end
+
 function Base.map!(f,r::MPIArray,args::MPIArray...)
     a = first(args)
     @assert all(i->size(a)==size(i),args)

src/p_range.jl

@@ -433,19 +433,18 @@ end
 
 function assembly_local_indices(indices,neighbors_snd,neighbors_rcv)
     cache = map(assembly_cache,indices)
-    mask =  map(cache) do cache
-        isassigned(cache.local_indices_snd) && isassigned(cache.local_indices_rcv)
+    mask =  map(cache) do mycache
+        isassigned(mycache.local_indices_snd) && isassigned(mycache.local_indices_rcv)
     end
     if ! getany(mask)
-        local_indices_snd, local_indices_rcv = compute_assembly_local_indices(indices,neighbors_snd,neighbors_rcv)
-        map(cache,local_indices_snd,local_indices_rcv) do cache, local_indices_snd, local_indices_rcv
-            cache.local_indices_snd[] = local_indices_snd
-            cache.local_indices_rcv[] = local_indices_rcv
+        new_local_indices_snd, new_local_indices_rcv = compute_assembly_local_indices(indices,neighbors_snd,neighbors_rcv)
+        map(cache,new_local_indices_snd,new_local_indices_rcv) do mycache, mylocal_indices_snd, mylocal_indices_rcv
+            mycache.local_indices_snd[] = mylocal_indices_snd
+            mycache.local_indices_rcv[] = mylocal_indices_rcv
         end
-        return local_indices_snd, local_indices_rcv
     end
-    local_indices_snd, local_indices_rcv = map(cache) do cache
-        cache.local_indices_snd[], cache.local_indices_rcv[]
+    local_indices_snd, local_indices_rcv = map(cache) do mycache
+        mycache.local_indices_snd[], mycache.local_indices_rcv[]
     end |> tuple_of_arrays
     local_indices_snd, local_indices_rcv
 end
@@ -476,20 +475,20 @@ function compute_assembly_local_indices(indices,neighbors_snd,neighbors_rcv)
             end
         end
         rewind_ptrs!(ptrs)
-        local_indices_snd = JaggedArray(data_lids,ptrs)
-        global_indices_snd = JaggedArray(data_gids,ptrs)
-        local_indices_snd, global_indices_snd
+        my_local_indices_snd = JaggedArray(data_lids,ptrs)
+        my_global_indices_snd = JaggedArray(data_gids,ptrs)
+        my_local_indices_snd, my_global_indices_snd
     end |>  tuple_of_arrays
     graph = ExchangeGraph(parts_snd,parts_rcv)
     global_indices_rcv = exchange_fetch(global_indices_snd,graph)
-    local_indices_rcv = map(global_indices_rcv,indices) do global_indices_rcv,indices
-        ptrs = global_indices_rcv.ptrs
+    local_indices_rcv = map(global_indices_rcv,indices) do myglobal_indices_rcv,myindices
+        ptrs = myglobal_indices_rcv.ptrs
         data_lids = zeros(Int32,ptrs[end]-1)
-        global_to_local_indices = global_to_local(indices)
-        for (k,gid) in enumerate(global_indices_rcv.data)
+        global_to_local_indices = global_to_local(myindices)
+        for (k,gid) in enumerate(myglobal_indices_rcv.data)
             data_lids[k] = global_to_local_indices[gid]
         end
-        local_indices_rcv = JaggedArray(data_lids,ptrs)
+        my_local_indices_rcv = JaggedArray(data_lids,ptrs)
     end
     local_indices_snd,local_indices_rcv
 end
@@ -568,9 +567,7 @@ function uniform_partition(rank,n::Integer,ghost::Bool,periodic::Bool=false)
     uniform_partition(rank,length(rank),n,ghost,periodic)
 end
 
-function uniform_partition(rank,np::Integer,n::Integer)
-    uniform_partition(rank,(np,),(n,))
-end
+function uniform_partition(rank,np::Integer,n::Integer) uniform_partition(rank,(np,),(n,)) end
 
 function uniform_partition(rank,np::Integer,n::Integer,ghost::Bool,periodic::Bool=false)
     uniform_partition(rank,(np,),(n,),(ghost,),(periodic,))
@@ -589,19 +586,19 @@ function block_with_constant_size(rank,np,n,ghost,periodic=map(i->false,ghost))
     own_ranges = map(local_range,Tuple(p),np,n)
     local_ranges = map(local_range,Tuple(p),np,n,ghost,periodic)
     owners = map(Tuple(p),own_ranges,local_ranges) do p,or,lr
-        owners = zeros(Int32,length(lr))
+        myowners = zeros(Int32,length(lr))
         for i in 1:length(lr)
             if lr[i] in or
-                owners[i] = p
+                myowners[i] = p
             end
         end
-        if owners[1] == 0
-            owners[1] = p-1
+        if myowners[1] == 0
+            myowners[1] = p-1
         end
-        if owners[end] == 0
-            owners[end] = p+1
+        if myowners[end] == 0
+            myowners[end] = p+1
         end
-        owners
+        myowners
     end
     n_ghost = 0
     cis = CartesianIndices(map(length,local_ranges))

src/p_sparse_matrix.jl

@@ -506,12 +506,13 @@ function to_trivial_partition(b::PVector,row_partition_in_main)
     destination = 1
     T = eltype(b)
     b_in_main = similar(b,T,PRange(row_partition_in_main))
-    map(own_values(b),partition(b_in_main),partition(axes(b,1))) do bown,b_in_main,indices
+    fill!(b_in_main,zero(T))
+    map(own_values(b),partition(b_in_main),partition(axes(b,1))) do bown,my_b_in_main,indices
         part = part_id(indices)
         if part == destination
-            b_in_main[own_to_global(indices)] .= bown
+            my_b_in_main[own_to_global(indices)] .= bown
         else
-            b_in_main .= bown
+            my_b_in_main .= bown
         end
     end
     assemble!(b_in_main) |> wait
@@ -521,12 +522,12 @@ end
 function from_trivial_partition!(c::PVector,c_in_main::PVector)
     destination = 1
     consistent!(c_in_main) |> wait
-    map(own_values(c),partition(c_in_main),partition(axes(c,1))) do cown, c_in_main, indices
+    map(own_values(c),partition(c_in_main),partition(axes(c,1))) do cown, my_c_in_main, indices
         part = part_id(indices)
         if part == destination
-            cown .= view(c_in_main,own_to_global(indices))
+            cown .= view(my_c_in_main,own_to_global(indices))
         else
-            cown .= c_in_main
+            cown .= my_c_in_main
         end
     end
     c
@@ -549,33 +550,33 @@ function to_trivial_partition(
                 n += 1
             end
         end
-        I = zeros(Int,n)
-        J = zeros(Int,n)
-        V = zeros(Ta,n)
+        myI = zeros(Int,n)
+        myJ = zeros(Int,n)
+        myV = zeros(Ta,n)
         n = 0
         for (i,j,v) in nziterator(a)
             if local_row_to_owner[i] == owner
                 n += 1
-                I[n] = local_to_global_row[i]
-                J[n] = local_to_global_col[j]
-                V[n] = v
+                myI[n] = local_to_global_row[i]
+                myJ[n] = local_to_global_col[j]
+                myV[n] = v
             end
         end
-        I,J,V
+        myI,myJ,myV
     end |> tuple_of_arrays
     assemble_coo!(I,J,V,row_partition_in_main) |> wait
-    I,J,V = map(partition(axes(a,1)),I,J,V) do row_indices,I,J,V
+    I,J,V = map(partition(axes(a,1)),I,J,V) do row_indices,myI,myJ,myV
         owner = part_id(row_indices)
         if owner == destination
-            I,J,V
+            myI,myJ,myV
         else
-            similar(I,eltype(I),0),similar(J,eltype(J),0),similar(V,eltype(V),0)
+            similar(myI,eltype(myI),0),similar(myJ,eltype(myJ),0),similar(myV,eltype(myV),0)
         end
     end |> tuple_of_arrays
-    values = map(I,J,V,row_partition_in_main,col_partition_in_main) do I,J,V,row_indices,col_indices
+    values = map(I,J,V,row_partition_in_main,col_partition_in_main) do myI,myJ,myV,row_indices,col_indices
         m = local_length(row_indices)
         n = local_length(col_indices)
-        compresscoo(M,I,J,V,m,n)
+        compresscoo(M,myI,myJ,myV,m,n)
     end
     PSparseMatrix(values,row_partition_in_main,col_partition_in_main)
 end
@@ -586,11 +587,12 @@ function Base.:\(a::PSparseMatrix,b::PVector)
     Tb = eltype(b)
     T = typeof(one(Ta)\one(Tb)+one(Ta)\one(Tb))
     c = PVector{Vector{T}}(undef,partition(axes(a,2)))
+    fill!(c,zero(T))
     a_in_main = to_trivial_partition(a)
     b_in_main = to_trivial_partition(b,partition(axes(a_in_main,1)))
     c_in_main = to_trivial_partition(c,partition(axes(a_in_main,2)))
-    map_main(partition(c_in_main),partition(a_in_main),partition(b_in_main)) do c, a, b
-        c .= a\b
+    map_main(partition(c_in_main),partition(a_in_main),partition(b_in_main)) do myc, mya, myb
+        myc .= mya\myb
         nothing
     end
     from_trivial_partition!(c,c_in_main)

src/p_vector.jl

@@ -251,10 +251,10 @@ function copy_cache(a::VectorAssemblyCache)
                     buffer_rcv)
 end
 function p_vector_cache_impl(::Type,vector_partition,index_partition)
-    neighbors = assembly_neighbors(index_partition)
-    indices = assembly_local_indices(index_partition,neighbors...)
-    buffers = map(assembly_buffers,vector_partition,indices...) |> tuple_of_arrays
-    map(VectorAssemblyCache,neighbors...,indices...,buffers...)
+    neighbors_snd,neighbors_rcv= assembly_neighbors(index_partition)
+    indices_snd,indices_rcv = assembly_local_indices(index_partition,neighbors_snd,neighbors_rcv)
+    buffers_snd,buffers_rcv = map(assembly_buffers,vector_partition,indices_snd,indices_rcv) |> tuple_of_arrays
+    map(VectorAssemblyCache,neighbors_snd,neighbors_rcv,indices_snd,indices_rcv,buffers_snd,buffers_rcv)
 end
 function assembly_buffers(values,local_indices_snd,local_indices_rcv)
     T = eltype(values)
@@ -267,8 +267,8 @@ function assembly_buffers(values,local_indices_snd,local_indices_rcv)
     buffer_snd, buffer_rcv
 end
 
-struct JaggedArrayAssemblyCache
-    cache::VectorAssemblyCache
+struct JaggedArrayAssemblyCache{T}
+    cache::VectorAssemblyCache{T}
 end
 Base.reverse(a::JaggedArrayAssemblyCache) = JaggedArrayAssemblyCache(reverse(a.cache))
 copy_cache(a::JaggedArrayAssemblyCache) = JaggedArrayAssemblyCache(copy_cache(a.cache))
@@ -518,7 +518,7 @@ Equivalent to
     PVector(vector_partition,index_partition)
 
 """
-function pvector(f,index_partition)
+@inline function pvector(f,index_partition)
     vector_partition = map(f,index_partition)
     PVector(vector_partition,index_partition)
 end