{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2133c064",
   "metadata": {},
   "source": [
    "<img src=\"https://upload.wikimedia.org/wikipedia/commons/thumb/3/39/VU_logo.png/800px-VU_logo.png?20161029201021\" width=\"350\">\n",
    "\n",
    "### Programming large-scale parallel systems\n",
    "### Lectures on Julia for HPC\n",
    "\n",
    "\n",
    "# Distributed Jacobi (SOR) method\n",
    "\n",
    "by Francesc Verdugo (VU Amsterdam)\n",
    "\n",
    "Version fall 2022"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7b64d5a",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- How to parallelize the Jacobi method in Julia"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bfa86d6",
   "metadata": {},
   "source": [
    "## Mathematical background\n",
    "\n",
    "\n",
    "### 1D Laplace equation\n",
    "\n",
    "\n",
    "Find a function $u(x)$ such that\n",
    "\n",
    "\n",
    "$u''(x) = 0 $ for $x\\in(0,L)$\n",
    "\n",
    "$u(0) = -1$\n",
    "\n",
    "$u(L) = 1$\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de45f04c",
   "metadata": {},
   "source": [
    "### Analytical solution\n",
    "\n",
    "$u(x) = (x-L)/L$"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bba2fb43",
   "metadata": {},
   "source": [
    "\n",
    " <div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 5724 characters]\" align=\"left\" width=\"250\"/>\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "392dd6e1",
   "metadata": {},
   "source": [
    "### A 3D Laplace equation in a complex domain"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e7d815e0",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 96028 characters]\" align=\"left\" width=\"500\"/>\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1450f45",
   "metadata": {},
   "source": [
    "### The jacobi method"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e794b036",
   "metadata": {},
   "source": [
    " <div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 15188 characters]\" align=\"left\" width=\"300\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bb20b40a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "jacobi (generic function with 1 method)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function jacobi(n,nsteps)\n",
    "    u = zeros(n+2)\n",
    "    u_new = zeros(n+2)\n",
    "    u[1] = -1\n",
    "    u[end] = 1\n",
    "    for istep in 1:nsteps\n",
    "        for i in 2:(n+1)\n",
    "            u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "        u[2:end-1] = u_new[2:end-1]\n",
    "    end\n",
    "    u\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9d251dba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7-element Vector{Float64}:\n",
       " -1.0\n",
       " -0.6666666666666666\n",
       " -0.3333333333333333\n",
       "  0.0\n",
       "  0.3333333333333333\n",
       "  0.6666666666666666\n",
       "  1.0"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jacobi(5,1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dda71761",
   "metadata": {},
   "source": [
    "## Parallelization strategy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c5a0d60",
   "metadata": {},
   "source": [
    "### Data partition and dependencies\n",
    "\n",
    "`u_new[i] = 0.5*(u[i-1]+u[i+1])`"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "52626d50",
   "metadata": {},
   "source": [
    " <div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 8924 characters]\" align=\"left\" width=\"450\"/>\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5e089f4",
   "metadata": {},
   "source": [
    "### \"Ghost\" entries"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1dc26426",
   "metadata": {},
   "source": [
    "\n",
    " <div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 12384 characters]\" align=\"left\" width=\"450\"/>\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acbbd646",
   "metadata": {},
   "source": [
    "### Remote channels"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "08dab820",
   "metadata": {},
   "source": [
    "\n",
    " <div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 24328 characters]\" align=\"left\" width=\"450\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "034fa536",
   "metadata": {},
   "source": [
    "## Implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a98e7ef1",
   "metadata": {},
   "source": [
    "### Disclaimer\n",
    "\n",
    "The most natural way of implementing a distributed Jacobi method is to use a library that provides distributed arrays:\n",
    "\n",
    "- [DistributedArrays.jl](https://github.com/JuliaParallel/DistributedArrays.jl)\n",
    "- [PartitionedArrays.jl](https://github.com/fverdugo/PartitionedArrays.jl)\n",
    "\n",
    "We will use low-level Julia functions as an academic exercise.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a2437ae",
   "metadata": {},
   "source": [
    "### Adding some procs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d95ff334",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Distributed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d64e580",
   "metadata": {},
   "outputs": [],
   "source": [
    "addprocs(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cd52702",
   "metadata": {},
   "source": [
    "### Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4aad98a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@everywhere function initialize_vectors(n)\n",
    "  @assert mod(n,nworkers()) == 0\n",
    "  local_n = div(n,nworkers())\n",
    "  u = zeros(local_n+2)\n",
    "  u_new = zeros(local_n+2)\n",
    "  u[1] = -1.0\n",
    "  u[end] = 1.0\n",
    "  u,u_new\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fceebbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "@everywhere function compute_error(u,n)\n",
    "  local_n = length(u)-2\n",
    "  u_exact = copy(u)\n",
    "  p = myid() - 1\n",
    "  for local_i in 0:(local_n+1)\n",
    "    global_i =  local_i + (p-1)*local_n + 1\n",
    "    Δ = 1/(n+1)\n",
    "    u_exact[local_i+1] = 2*(Δ*global_i-Δ)-1\n",
    "  end\n",
    "  maximum(abs.(u[2:end-1].-u_exact[2:end-1]))\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd06d27f",
   "metadata": {},
   "source": [
    "### Distributed Jacobi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56d3bd49",
   "metadata": {},
   "outputs": [],
   "source": [
    "function jacobi_dist(n,nsteps)\n",
    "  nw = nworkers()\n",
    "  ftrs_prev_snd = Vector{Future}(undef,nw)\n",
    "  ftrs_next_snd = Vector{Future}(undef,nw)\n",
    "  fun = ()->Channel{Float64}(1)\n",
    "  for w in workers()\n",
    "    p = w-1\n",
    "    ftrs_prev_snd[p] = @spawnat w RemoteChannel(fun,myid())\n",
    "    ftrs_next_snd[p] = @spawnat w RemoteChannel(fun,myid())\n",
    "  end\n",
    "  ftrs = Vector{Future}(undef,nw)\n",
    "  @sync for w in workers()\n",
    "    ftrs[w-1] = @spawnat w jacobi_worker(\n",
    "            n,nsteps,ftrs_prev_snd,ftrs_next_snd)\n",
    "  end\n",
    "  reduce(max,fetch.(ftrs))\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab2c25f3",
   "metadata": {},
   "source": [
    "### Exercise: complete the cell below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "739ccfa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "@everywhere function jacobi_worker(\n",
    "  n,nsteps,ftrs_prev_snd,ftrs_next_snd)\n",
    "    u, u_new = initialize_vectors(n)\n",
    "    w = myid(); p = w-1\n",
    "    # TODO: Setup channels\n",
    "    #chn_prev_snd\n",
    "    #chn_prev_rcv\n",
    "    #chn_next_snd\n",
    "    #chn_next_rcv\n",
    "    for step in 1:nsteps       \n",
    "        # TODO: Communicate\n",
    "        # Local computations\n",
    "        for i in 2:(length(u)-1)\n",
    "            u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "        u[2:end-1] = u_new[2:end-1]\n",
    "    end\n",
    "    @show u\n",
    "    err = compute_error(u,n)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edbb2318",
   "metadata": {},
   "source": [
    "### Test-driven development"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2301383",
   "metadata": {},
   "outputs": [],
   "source": [
    "err = jacobi_dist(6,1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63ec8554",
   "metadata": {},
   "source": [
    "### Remember: remote channels"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d4f59c15",
   "metadata": {},
   "source": [
    " <div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 24328 characters]\" align=\"left\" width=\"450\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2d8bdbc",
   "metadata": {},
   "source": [
    "## Possible enhancements"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0af07603",
   "metadata": {},
   "source": [
    "### Is it possible to overlap computation and communication?"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d401c647",
   "metadata": {},
   "source": [
    "\n",
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 12384 characters]\" align=\"left\" width=\"450\"/>\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80e67eb9",
   "metadata": {},
   "source": [
    "### 2nd exercise:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9aaa50c",
   "metadata": {},
   "outputs": [],
   "source": [
    "@everywhere function jacobi_worker(\n",
    "  n,nsteps,ftrs_prev_snd,ftrs_next_snd)\n",
    "    u, u_new = initialize_vectors(n)\n",
    "    p = myid()-1\n",
    "    chn_prev_snd = fetch(ftrs_prev_snd[p])\n",
    "    chn_next_snd = fetch(ftrs_next_snd[p])\n",
    "    if myid()!=2\n",
    "        chn_prev_rcv = fetch(ftrs_next_snd[p-1])\n",
    "    end\n",
    "    if myid() != nprocs()\n",
    "        chn_next_rcv = fetch(ftrs_prev_snd[p+1])\n",
    "    end\n",
    "    for step in 1:nsteps\n",
    "        # TODO overlap communication and computation\n",
    "        if myid() != 2\n",
    "           put!(chn_prev_snd,u[2])\n",
    "           u[1] = take!(chn_prev_rcv)\n",
    "        end\n",
    "        if myid() != nprocs()\n",
    "           put!(chn_next_snd,u[end-1])\n",
    "           u[end] = take!(chn_next_rcv)\n",
    "        end\n",
    "        for i in 2:(length(u)-1)\n",
    "            u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "        u[2:end-1] = u_new[2:end-1]\n",
    "    end\n",
    "    @show u\n",
    "    err = compute_error(u,n)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9f12348",
   "metadata": {},
   "source": [
    "## Solution of the exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7d067f8",
   "metadata": {},
   "source": [
    "### 1st exercise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e2514b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "@everywhere function jacobi_worker(\n",
    "  n,nsteps,ftrs_prev_snd,ftrs_next_snd)\n",
    "    u, u_new = initialize_vectors(n)\n",
    "    p = myid()-1\n",
    "    chn_prev_snd = fetch(ftrs_prev_snd[p])\n",
    "    chn_next_snd = fetch(ftrs_next_snd[p])\n",
    "    if myid()!=2\n",
    "        chn_prev_rcv = fetch(ftrs_next_snd[p-1])\n",
    "    end\n",
    "    if myid() != nprocs()\n",
    "        chn_next_rcv = fetch(ftrs_prev_snd[p+1])\n",
    "    end\n",
    "    for step in 1:nsteps       \n",
    "        if myid() != 2\n",
    "           put!(chn_prev_snd,u[2])\n",
    "           u[1] = take!(chn_prev_rcv)\n",
    "        end\n",
    "        if myid() != nprocs()\n",
    "           put!(chn_next_snd,u[end-1])\n",
    "           u[end] = take!(chn_next_rcv)\n",
    "        end\n",
    "        # Local computations\n",
    "        for i in 2:(length(u)-1)\n",
    "            u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "        u[2:end-1] = u_new[2:end-1]\n",
    "    end\n",
    "    @show u\n",
    "    err = compute_error(u,n)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "157e705d",
   "metadata": {},
   "source": [
    "### 2nd exercise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6f5a89f",
   "metadata": {},
   "outputs": [],
   "source": [
    "@everywhere function jacobi_worker(\n",
    "  n,nsteps,ftrs_prev_snd,ftrs_next_snd)\n",
    "    u, u_new = initialize_vectors(n)\n",
    "    p = myid()-1\n",
    "    chn_prev_snd = fetch(ftrs_prev_snd[p])\n",
    "    chn_next_snd = fetch(ftrs_next_snd[p])\n",
    "    if myid()!=2\n",
    "        chn_prev_rcv = fetch(ftrs_next_snd[p-1])\n",
    "    end\n",
    "    if myid() != nprocs()\n",
    "        chn_next_rcv = fetch(ftrs_prev_snd[p+1])\n",
    "    end\n",
    "    for step in 1:nsteps       \n",
    "        if myid() != 2\n",
    "           put!(chn_prev_snd,u[2])\n",
    "           t1 = @async u[1] = take!(chn_prev_rcv)\n",
    "        end\n",
    "        if myid() != nprocs()\n",
    "           put!(chn_next_snd,u[end-1])\n",
    "           t2 = @async u[end] = take!(chn_next_rcv)\n",
    "        end\n",
    "        # Local computations (interior)\n",
    "        for i in 3:(length(u)-2)\n",
    "            u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "        # Wait for ghost values\n",
    "        if myid() != 2\n",
    "            wait(t1)\n",
    "        end\n",
    "        if myid() != nprocs()\n",
    "            wait(t2)\n",
    "        end\n",
    "        # Update near boundary values\n",
    "        for i in [2,length(u)-1]\n",
    "            u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "        u[2:end-1] = u_new[2:end-1]\n",
    "    end\n",
    "    @show u\n",
    "    err = compute_error(u,n)\n",
    "end"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.0",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
