Enhancements in Jacobi notebook

docs/make.jl

@@ -122,7 +122,7 @@ makedocs(;
                          #"Distributed computing with MPI" => "mpi_tutorial.md",
                          "Matrix-matrix multiplication"=>"matrix_matrix.md",
                          "MPI (point-to-point)" => "julia_mpi.md",
-                         #"Jacobi method" => "jacobi_method.md",
+                         "Jacobi method" => "jacobi_method.md",
                          #"All pairs of shortest paths" => "asp.md",
                          #"Gaussian elimination" => "LEQ.md",
                          #"Traveling salesperson problem" => "tsp.md",

docs/src/solutions_for_all_notebooks.md

@@ -167,50 +167,70 @@ end
 ### Exercise 1
 
 ```julia
-@everywhere workers() begin
-    using MPI
-    comm = MPI.Comm_dup(MPI.COMM_WORLD)
-    function jacobi_mpi(n,niters)
-        nranks = MPI.Comm_size(comm)
-        rank = MPI.Comm_rank(comm)
-        if mod(n,nranks) != 0
-            println("n must be a multiple of nranks")
-            MPI.Abort(comm,1)
-        end
-        n_own = div(n,nranks)
-        u = zeros(n_own+2)
-        u[1] = -1
-        u[end] = 1
-        u_new = copy(u)
-        for t in 1:niters
-            reqs = MPI.Request[]
-            if rank != 0
-                neig_rank = rank-1
-                req = MPI.Isend(view(u,2:2),comm,dest=neig_rank,tag=0)
-                push!(reqs,req)
-                req = MPI.Irecv!(view(u,1:1),comm,source=neig_rank,tag=0)
-                push!(reqs,req)
-            end
-            if rank != (nranks-1)
-                neig_rank = rank+1
-                s = n_own+1
-                r = n_own+2
-                req = MPI.Isend(view(u,s:s),comm,dest=neig_rank,tag=0)
-                push!(reqs,req)
-                req = MPI.Irecv!(view(u,r:r),comm,source=neig_rank,tag=0)
-                push!(reqs,req)
-            end
-            for i in 3:n_own
-                u_new[i] = 0.5*(u[i-1]+u[i+1])
-            end
-            MPI.Waitall(reqs)
-            for i in (2,n_own+1)
-                u_new[i] = 0.5*(u[i-1]+u[i+1])
-            end
-            u, u_new = u_new, u
-        end
-        return u
+function jacobi_mpi(n,niters)
+    comm = MPI.COMM_WORLD
+    nranks = MPI.Comm_size(comm)
+    rank = MPI.Comm_rank(comm)
+    if mod(n,nranks) != 0
+        println("n must be a multiple of nranks")
+        MPI.Abort(comm,1)
     end
+    load = div(n,nranks)
+    u = zeros(load+2)
+    u[1] = -1
+    u[end] = 1
+    u_new = copy(u)
+    for t in 1:niters
+        reqs = MPI.Request[]
+        if rank != 0
+            neig_rank = rank-1
+            req = MPI.Isend(view(u,2:2),comm,dest=neig_rank,tag=0)
+            push!(reqs,req)
+            req = MPI.Irecv!(view(u,1:1),comm,source=neig_rank,tag=0)
+            push!(reqs,req)
+        end
+        if rank != (nranks-1)
+            neig_rank = rank+1
+            s = load+1
+            r = load+2
+            req = MPI.Isend(view(u,s:s),comm,dest=neig_rank,tag=0)
+            push!(reqs,req)
+            req = MPI.Irecv!(view(u,r:r),comm,source=neig_rank,tag=0)
+            push!(reqs,req)
+        end
+        for i in 3:load
+            u_new[i] = 0.5*(u[i-1]+u[i+1])
+        end
+        MPI.Waitall(reqs)
+        for i in (2,load+1)
+            u_new[i] = 0.5*(u[i-1]+u[i+1])
+        end
+        u, u_new = u_new, u
+
+    end
+    # Gather the results
+    if rank !=0
+        lb = 2
+        ub = load+1
+        MPI.Send(view(u,lb:ub),comm,dest=0)
+        u_all = zeros(0) # This will nevel be used
+    else
+        u_all = zeros(n+2)
+        # Set boundary
+        u_all[1] = -1
+        u_all[end] = 1
+        # Set data for rank 0
+        lb = 2
+        ub = load+1
+        u_all[lb:ub] = view(u,lb:ub)
+        # Set data for other ranks
+        for other_rank in 1:(nranks-1)
+            lb += load
+            ub += load
+            MPI.Recv!(view(u_all,lb:ub),comm;source=other_rank)
+        end
+    end
+    return u_all
 end
 ```