parallel-gpu-workshop-JuliaCon21/scripts/diffusion_1D_mpi.jl at main · luraess/parallel-gpu-workshop-JuliaCon21 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
# 1D linear diffusion Julia MPI solver
# run: ~/.julia/bin/mpiexecjl -n 4 julia --project diffusion_1D_mpi.jl
using Plots, Printf, MAT
import MPI

# enable plotting by default
if !@isdefined do_save; do_save = true end

# MPI functions
@views function update_halo(A, neighbors_x, comm)
    if neighbors_x[1] != MPI.MPI_PROC_NULL
        sendbuf = A[2]
        recvbuf = zeros(1)
        MPI.Send(sendbuf,  neighbors_x[1], 0, comm)
        MPI.Recv!(recvbuf, neighbors_x[1], 1, comm)
        A[1] = recvbuf[1]
    end
    if neighbors_x[2] != MPI.MPI_PROC_NULL
        sendbuf = A[end-1]
        recvbuf = zeros(1)
        MPI.Send(sendbuf,  neighbors_x[2], 1, comm)
        MPI.Recv!(recvbuf, neighbors_x[2], 0, comm)
        A[end] = recvbuf[1]
    end
    return
end

@views function diffusion_1D_mpi(; do_save=false)
    # MPI
    MPI.Init()
    dims        = [0]
    comm        = MPI.COMM_WORLD
    nprocs      = MPI.Comm_size(comm)
    MPI.Dims_create!(nprocs, dims)
    comm_cart   = MPI.Cart_create(comm, dims, [0], 1)
    me          = MPI.Comm_rank(comm_cart)
    coords      = MPI.Cart_coords(comm_cart)
    neighbors_x = MPI.Cart_shift(comm_cart, 0, 1)
    if (me==0) println("nprocs=$(nprocs), dims[1]=$(dims[1])") end
    # Physics
    lx    = 10.0
    λ     = 1.0
    nt    = 100
    # Numerics
    nx    = 32                 # local number of grid points
    nx_g  = dims[1]*(nx-2) + 2 # global number of grid points
    # Derived numerics
    dx    = lx/nx_g            # global
    dt    = dx^2/λ/2.1
    # Array allocation
    qHx   = zeros(nx-1)
    # Initial condition
    x0    = coords[1]*(nx-2)*dx
    xc    = [x0 + ix*dx - dx/2 - 0.5*lx  for ix=1:nx]
    H     = exp.(.-xc.^2)
    t_tic = 0.0
    # Time loop
    for it = 1:nt
        if (it==11) t_tic = Base.time() end
        qHx        .= .-λ*diff(H)/dx
        H[2:end-1] .= H[2:end-1] .- dt*diff(qHx)/dx
        update_halo(H, neighbors_x, comm_cart)
    end
    t_toc = Base.time()-t_tic
    if (me==0) @printf("Time = %1.4e s, T_eff = %1.2f GB/s \n", t_toc, round((2/1e9*nx*sizeof(lx))/(t_toc/(nt-10)), sigdigits=2)) end
    if do_save file = matopen("$(@__DIR__)/H_$(me).mat", "w"); write(file, "H", Array(H)); close(file) end
    MPI.Finalize()
    return
end

diffusion_1D_mpi(; do_save=do_save)