Partition of distributed arrays in Julia

General Usage Performance
1

0

my problem is to complete partition_data(strain::AbstractArray, stress::AbstractArray) in Julia. In this function, the strain/stress arguments are Arrays that contain the entire stress and strain data only on the processor running the Julia REPL. This function should distribute the strain/stress data to all processors available by turning them into DistributedArrays.

I have the code using MPI, how can I convert it to array distribution? I tried to convert the strain and stress arrays to distributed arrays but I’m confused how to do the partition.

Blockquote

function partition_data(strain::AbstractArray, stress::AbstractArray, comm::MPI.Comm)

    rank = MPI.Comm_rank(comm)
    size = MPI.Comm_size(comm)

    len = length(strain)
    lengths = nothing
    offset = nothing

    if rank == 0
        remainder = len % size
        chunk_size = len ÷ size
        lengths = ones(Int, size) .* chunk_size
        lengths[1:remainder] .*= 1
        lengths[2:end] .+= 1
        offset = zeros(Int, size)
        offset[2:end] = cumsum(lengths)[1:end-1] .- collect(1:size-1)
    end    

    lengths = MPI.bcast(lengths, 0, comm)
    offset = MPI.bcast(offset, 0, comm)

    my_stress = Vector{Float64}(undef, lengths[rank+1])
    my_strain = Vector{Float64}(undef, lengths[rank+1])

    MPI.Scatterv!(MPI.VBuffer(strain, lengths, offset, MPI.DOUBLE), my_strain, 0, comm)
    MPI.Scatterv!(MPI.VBuffer(stress, lengths, offset, MPI.DOUBLE), my_stress, 0, comm)
    (my_strain, my_stress)
end