Manifest.toml

# This file is machine-generated - editing it directly is not advised

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[CompilerSupportLibraries_jll]]
deps = ["Libdl", "Pkg"]
git-tree-sha1 = "7c4f882c41faa72118841185afc58a2eb00ef612"
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
version = "0.3.3+0"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[DocStringExtensions]]
deps = ["LibGit2", "Markdown", "Pkg", "Test"]
git-tree-sha1 = "c5714d9bcdba66389612dc4c47ed827c64112997"
uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
version = "0.8.2"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[LibGit2]]
deps = ["Printf"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[MPI]]
deps = ["Distributed", "DocStringExtensions", "Libdl", "MPICH_jll", "MicrosoftMPI_jll", "OpenMPI_jll", "Pkg", "Random", "Requires", "Serialization", "Sockets"]
git-tree-sha1 = "925f4e00b51d7a68ecb88c1b60d3b7aefb8c8fcb"
uuid = "da04e1cc-30fd-572f-bb4f-1f8673147195"
version = "0.15.0"

[[MPICH_jll]]
deps = ["CompilerSupportLibraries_jll", "Libdl", "Pkg"]
git-tree-sha1 = "4d37f1e07b4e2a74462eebf9ee48c626d15ffdac"
uuid = "7cb0a576-ebde-5e09-9194-50597f1243b4"
version = "3.3.2+10"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[MicrosoftMPI_jll]]
deps = ["Libdl", "Pkg"]
git-tree-sha1 = "720de13004e416f2864c92593d3839e062703985"
uuid = "9237b28f-5490-5468-be7b-bb81f5f5e6cf"
version = "10.1.2+3"

[[OpenMPI_jll]]
deps = ["Libdl", "Pkg"]
git-tree-sha1 = "41b983e26a7ab8c9bf05f7d70c274b817d541b46"
uuid = "fe0851c0-eecd-5654-98d4-656369965a5c"
version = "4.0.2+2"

[[Pkg]]
deps = ["Dates", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "UUIDs"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[Requires]]
deps = ["UUIDs"]
git-tree-sha1 = "d37400976e98018ee840e0ca4f9d20baa231dc6b"
uuid = "ae029012-a4dd-5104-9daa-d747884805df"
version = "1.0.1"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[Test]]
deps = ["Distributed", "InteractiveUtils", "Logging", "Random"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

mpi-variance-array.jl

using Statistics, MPI, Printf

MPI.Init()


struct SummaryStat
    mean::Float64
    var::Float64
    n::Float64
end

function SummaryStat(X::Vector)
    m = mean(X)
    v = varm(X,m, corrected=false)
    n = length(X)
    SummaryStat(m,v,n)
end

function combine(S1::SummaryStat, S2::SummaryStat)
    n = S1.n + S2.n
    m = (S1.mean*S1.n + S2.mean*S2.n) / n

    v = (S1.n * (S1.var + S1.mean * (S1.mean-m)) + S2.n * (S2.var + S2.mean * (S2.mean-m)))/n

    SummaryStat(m,v,n)
end


function main(N)
    comm = MPI.COMM_WORLD
    isroot = MPI.Comm_rank(comm) == 0
    
    X = rand(100,100)

    # check we're computing the correct result
    recv = MPI.Reduce(mapslices(SummaryStat, X, dims=1), combine, 0, comm)
    Y = isroot ? zeros(100,100,MPI.Comm_size(comm)) : nothing
    MPI.Gather!(X, Y, 0, comm)
    if isroot
        @assert [r.var for r in recv] ≈ var(Y; dims=(1,3), corrected=false)    
    end
    
    T = zeros(N)
    MPI.Barrier(comm)
    
    for i = 1:N
        T[i] = @elapsed(recv = MPI.Reduce(mapslices(SummaryStat, X, dims=1), combine, 0, comm))
        MPI.Barrier(comm)
    end
    if MPI.Comm_rank(comm) == 0
        @printf "nranks: %4d, niters: %4d, time (min/avg/max): %10.8f / %10.8f/ %10.8f\n" MPI.Comm_size(comm) N minimum(T) mean(T) maximum(T)
    end
end

main(100)

mpi-variance.jl

using Statistics, MPI, Printf

MPI.Init()


struct SummaryStat
    mean::Float64
    var::Float64
    n::Float64
end

function SummaryStat(X::Vector)
    m = mean(X)
    v = varm(X,m, corrected=false)
    n = length(X)
    SummaryStat(m,v,n)
end

function combine(S1::SummaryStat, S2::SummaryStat)
    n = S1.n + S2.n
    m = (S1.mean*S1.n + S2.mean*S2.n) / n

    v = (S1.n * (S1.var + S1.mean * (S1.mean-m)) + S2.n * (S2.var + S2.mean * (S2.mean-m)))/n

    SummaryStat(m,v,n)
end


function main(N)
    comm = MPI.COMM_WORLD
    
    X = randn(100)

    recv = MPI.Reduce(SummaryStat(X), combine, 0, MPI.COMM_WORLD)

    T = zeros(N)
    MPI.Barrier(comm)
    
    for i = 1:N
        T[i] = @elapsed(recv = MPI.Reduce(SummaryStat(X), combine, 0, MPI.COMM_WORLD))
        MPI.Barrier(comm)
    end
    if MPI.Comm_rank(comm) == 0
        @printf "nranks: %4d, niters: %4d, time (min/avg/max): %10.8f / %10.8f/ %10.8f\n" MPI.Comm_size(comm) N minimum(T) mean(T) maximum(T)
    end
end

main(100)

Project.toml

[deps]
MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

run.sh

#!/bin/bash

#SBATCH --time=00:10:00

module load openmpi/4.0.3 julia/1.4.2

julia --project -e 'using Pkg; pkg"instantiate"; pkg"precompile"'

for ((n=1;n<=SLURM_JOB_NUM_NODES;n=n*2)); do
    mpiexec -n $n julia --project mpi-variance-array.jl
done