{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2e15ced8",
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8549215",
   "metadata": {},
   "source": [
    "# Jacobi method"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e0ef563",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "In this notebook, we will learn\n",
    "\n",
    "- How to paralleize the Jacobi method\n",
    "- How the data partition can impact the performance of a distributed algorithm\n",
    "- How to use latency hiding to improve parallel performance\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48fd29d9",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b> Do not forget to run the next cell before you start studying this notebook. \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dc78750",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Printf\n",
    "\n",
    "function answer_checker(answer,solution)\n",
    "    if answer == solution\n",
    "        \"🥳 Well done! \"\n",
    "    else\n",
    "        \"It's not correct. Keep trying! 💪\"\n",
    "    end |> println\n",
    "end\n",
    "gauss_seidel_1_check(answer) = answer_checker(answer,\"c\")\n",
    "jacobi_1_check(answer) = answer_checker(answer, \"d\")\n",
    "jacobi_2_check(answer) = answer_checker(answer, \"b\")\n",
    "jacobi_3_check(answer) = answer_checker(answer, \"c\")\n",
    "lh_check(answer) = answer_checker(answer, \"c\")\n",
    "sndrcv_check(answer) = answer_checker(answer,\"b\")\n",
    "function partition_1d_answer(bool)\n",
    "    bool || return\n",
    "    msg = \"\"\"\n",
    "- We update N^2/P items per iteration\n",
    "- We need data from 2 neighbors (2 messages per iteration)\n",
    "- We communicate N items per message\n",
    "- Communication/computation ratio is 2N/(N^2/P) = 2P/N =O(P/N)\n",
    "    \"\"\"\n",
    "    println(msg)\n",
    "end\n",
    "function partition_2d_answer(bool)\n",
    "        bool || return\n",
    "    msg = \"\"\"\n",
    "- We update N^2/P items per iteration\n",
    "- We need data from 4 neighbors (4 messages per iteration)\n",
    "- We communicate N/sqrt(P) items per message\n",
    "- Communication/computation ratio is (4N/sqrt(P)/(N^2/P)= 4sqrt(P)/N =O(sqrt(P)/N)\n",
    "    \"\"\"\n",
    "    println(msg)\n",
    "end\n",
    "function partition_cyclic_answer(bool)\n",
    "        bool || return\n",
    "    msg = \"\"\"\n",
    "- We update N^2/P items\n",
    "- We need data from 4 neighbors (4 messages per iteration)\n",
    "- We communicate N^2/P items per message (the full data owned by the neighbor)\n",
    "- Communication/computation ratio is O(1)\n",
    "    \"\"\"\n",
    "println(msg)\n",
    "end\n",
    "function sndrcv_fix_answer(bool)\n",
    "        bool || return\n",
    "    msg = \"\"\"\n",
    "    One needs to carefully order the sends and the receives to avoid cyclic dependencies\n",
    "    that might result in deadlocks. The actual implementation is left as an exercise.   \n",
    "    \"\"\"\n",
    "    println(msg)\n",
    "end\n",
    "jacobitest_check(answer) = answer_checker(answer,\"a\")\n",
    "function jacobitest_why(bool)\n",
    "        bool || return\n",
    "    msg = \"\"\"\n",
    "    The test will pass. The parallel implementation does exactly the same operations\n",
    "    in exactly the same order than the sequential one. Thus, the result should be\n",
    "    exactly the same. Note however this is often not true in other parallel algorithms.\n",
    "    Specially, when one does parallel reductions.\n",
    "    \"\"\"\n",
    "    println(msg)\n",
    "end\n",
    "gauss_seidel_2_check(answer) = answer_checker(answer,\"d\")\n",
    "function gauss_seidel_2_why(bool)\n",
    "        bool || return\n",
    "    msg = \"\"\"\n",
    "    All \"red\" cells can be updated in parallel as they only depend on the values of \"black\" cells.\n",
    "    In order workds, we can update the \"red\" cells in any order whithout changing the result. They only\n",
    "    depend on values in the \"black\" cells, which will not change during the loop over \"red\" cells.\n",
    "    Similarly, all \"black\" cells can be updated in parallel as they only depend on \"red\" cells.\n",
    "    \"\"\"\n",
    "    println(msg)\n",
    "end\n",
    "println(\"🥳 Well done! \")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4cb59d5",
   "metadata": {},
   "source": [
    "## The Jacobi method for the Laplace equation\n",
    "\n",
    "\n",
    "The [Jacobi method](https://en.wikipedia.org/wiki/Jacobi_method) is a numerical tool to solve systems of linear algebraic equations. One of the main applications of the Jacobi method is to solve the equations resulting from boundary value problems (BVPs). I.e., given the values at the boundary (of a grid), we are interested in finding the interior values that fulfill a certain equation.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a2bdbc6",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 15188 characters]\" align=\"left\" width=\"300\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93e84ff8",
   "metadata": {},
   "source": [
    "When solving a [Laplace equation](https://en.wikipedia.org/wiki/Laplace%27s_equation) in 1D, the Jacobi method leads to the following iterative scheme: The entry $i$ of vector $u$ at iteration $t+1$ is computed as:\n",
    "\n",
    "$u^{t+1}_i = \\dfrac{u^t_{i-1}+u^t_{i+1}}{2}$\n",
    "\n",
    "\n",
    "This algorithm is yet simple but shares fundamental challenges with many other algorithms used in scientific computing. This is why we are studying it here.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e63a5792",
   "metadata": {},
   "source": [
    "### Serial implementation\n",
    "\n",
    "The following code implements the iterative scheme above for boundary conditions -1 and 1 on a grid with $n$ interior points.\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b> `u, u_new = u_new, u` is equivalent to `tmp = u; u = u_new; u_new = tmp`. I.e. we swap the arrays `u` and `u_new` are referring to. \n",
    "</div>\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14a58308",
   "metadata": {},
   "outputs": [],
   "source": [
    "function jacobi(n,niters)\n",
    "    u = zeros(n+2)\n",
    "    u[1] = -1\n",
    "    u[end] = 1\n",
    "    u_new = copy(u)\n",
    "    for t in 1:niters\n",
    "        for i in 2:(n+1)\n",
    "            u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "        u, u_new = u_new, u\n",
    "    end\n",
    "    u\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "432bd862",
   "metadata": {},
   "source": [
    "If you run it for zero iterations, we will see the initial condition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76e1eba1",
   "metadata": {},
   "outputs": [],
   "source": [
    "jacobi(5,0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c75cb9a6",
   "metadata": {},
   "source": [
    "If you run it for enough iterations, you will see the expected solution of the Laplace equation: values that vary linearly form -1 to 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b52be374",
   "metadata": {},
   "outputs": [],
   "source": [
    "jacobi(5,100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22fda724",
   "metadata": {},
   "source": [
    "In our version of the Jacobi method, we return after a given number of iterations. Other stopping criteria are possible. For instance, iterate until the maximum difference between u and u_new (in absolute value) is below a tolerance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15de7bf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "function jacobi_with_tol(n,tol)\n",
    "    u = zeros(n+2)\n",
    "    u[1] = -1\n",
    "    u[end] = 1\n",
    "    u_new = copy(u)\n",
    "    while true\n",
    "        diff = 0.0\n",
    "        for i in 2:(n+1)\n",
    "            ui_new = 0.5*(u[i-1]+u[i+1])\n",
    "            u_new[i] = ui_new\n",
    "            diff_i = abs(ui_new-u[i])\n",
    "            diff = max(diff_i,diff)            \n",
    "        end\n",
    "        if diff < tol\n",
    "            return u_new\n",
    "        end\n",
    "        u, u_new = u_new, u\n",
    "    end\n",
    "    u\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "697ad307",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5\n",
    "tol = 1e-9\n",
    "jacobi_with_tol(n,tol)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e085701",
   "metadata": {},
   "source": [
    "However, we are not going to parallelize this more complex in this notebook (left as an exercise). The simpler one is already challenging enough to start with."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9df06442",
   "metadata": {},
   "source": [
    "## Parallelization of the Jacobi method\n",
    "\n",
    "Now, let us parallelize the Jacobi method."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6918c31",
   "metadata": {},
   "source": [
    "\n",
    "### Where can we exploit parallelism?\n",
    "\n",
    "Look at the two nested loops in the sequential implementation:\n",
    "\n",
    "```julia\n",
    "for t 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, u_new = u_new, u\n",
    "end\n",
    "```\n",
    "\n",
    "- The outer loop over `t` cannot be parallelized. The value of `u` at step `t+1` depends on the value at the previous step `t`.\n",
    "- The inner loop is trivially parallel. The loop iterations are independent (any order is possible).\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97a2d7d5",
   "metadata": {},
   "source": [
    "### Parallelization strategy\n",
    "\n",
    "Remember that a sufficiently large grain size is needed to achieve performance in a distributed algorithm. For Jacobi, one could update each entry of vector `u_new` in a different process, but this would not be efficient. Instead, we use a parallelization strategy with a larger grain size that is analogous to the algorithm 3 we studied for the matrix-matrix multiplication:\n",
    "\n",
    "- Data partition: each worker updates a consecutive section of the array `u_new` \n",
    "\n",
    "The following figure displays the data distribution over 3 processes."
   ]
  },
  {
   "attachments": {
    "fig09.png": {
     "image/png": "[encoded data, 19372 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "a9667c42",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 19372 characters]\"  align=\"left\" width=\"450\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22dc6e54",
   "metadata": {},
   "source": [
    "### Data dependencies\n",
    "\n",
    "Recall the Jacobi update:\n",
    "\n",
    "`u_new[i] = 0.5*(u[i-1]+u[i+1])`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba4113af",
   "metadata": {},
   "source": [
    "Note that an entry in the interior of the locally stored vector can be updated using local data  only. For updating this one, communication is not needed."
   ]
  },
  {
   "attachments": {
    "f1.png": {
     "image/png": "[encoded data, 8436 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "97a5079d",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 8436 characters]\"  align=\"left\" width=\"450\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "343e60c3",
   "metadata": {},
   "source": [
    "However, to update the entries on the boundary of the locally stored vector we need entries stored on other processors."
   ]
  },
  {
   "attachments": {
    "f2.png": {
     "image/png": "[encoded data, 8436 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "fce954fe",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 8436 characters]\"  align=\"left\" width=\"450\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f90d701",
   "metadata": {},
   "source": [
    "Thus, in order to update the local entries in `u_new`, we also need  some remote entries of vector `u` located in neighboring processes. Figure below shows the entries of `u` needed to update the local entries of `u_new` in a particular process (CPU 2)."
   ]
  },
  {
   "attachments": {
    "fig22.png": {
     "image/png": "[encoded data, 15076 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "8e156e29",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 15076 characters]\"  align=\"left\" width=\"450\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "513f1f7e",
   "metadata": {},
   "source": [
    "### Communication overhead\n",
    "\n",
    "Now that we understand which are the data dependencies, we can do the theoretical performance analysis.\n",
    "\n",
    "\n",
    "- We update $N/P$ entries in each process at each iteration, where $N$ is the total length of the vector and $P$ the number of processes\n",
    "- Thus, computation complexity is $O(N/P)$\n",
    "- We need to get remote entries from 2 neighbors (2 messages per iteration)\n",
    "- We need to communicate 1 entry per message\n",
    "- Thus, communication complexity is $O(1)$\n",
    "- Communication/computation ration is $O(P/N)$, making the algorithm potentially scalable if $P<<N$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b3c8c05",
   "metadata": {},
   "source": [
    "### Ghost (aka halo) cells\n",
    "\n",
    "This parallel strategy is efficient according to the theoretical analysis. But how to implement it? A usual way of implementing the Jacobi method and related algorithms is using so-called ghost cells. Ghost cells represent the missing data dependencies in the data owned by each process. After importing the appropriate values from the neighbor processes one can perform the usual sequential Jacobi update locally in the processes. Cells with gray edges are ghost (or boundary) cells in the following figure. Note that we added one ghost cell at the front and end of the local array."
   ]
  },
  {
   "attachments": {
    "fig13.png": {
     "image/png": "[encoded data, 12128 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "98e0eb71",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 12128 characters]\"  align=\"left\" width=\"450\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0148f9b3",
   "metadata": {},
   "source": [
    "Thus, the algorithm is usually implemented following two main phases at each iteration Jacobi:\n",
    "\n",
    "1. Fill the ghost entries with communications\n",
    "2. Do the Jacobi update sequentially at each process"
   ]
  },
  {
   "attachments": {
    "fig15.png": {
     "image/png": "[encoded data, 29440 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "baccd833",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 29440 characters]\"  align=\"left\" width=\"450\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a40846c",
   "metadata": {},
   "source": [
    "We are going to implement this algorithm with MPI later in this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75f735a2",
   "metadata": {},
   "source": [
    "## Extension to 2D\n",
    "\n",
    "\n",
    "The Jacobi method studied so far was for a one dimensional Laplace equation. In real-world applications however, one solve equations in multiple dimensions. Typically 2D and 3D. The 2D and 3D cases are conceptually equivalent, but we will discuss the 2D case here for simplicity.\n",
    "\n",
    "Now, the goal is to find the interior points of a 2D grid given the values at the boundary.\n",
    "\n"
   ]
  },
  {
   "attachments": {
    "fig17.png": {
     "image/png": "[encoded data, 45660 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "18659ed7",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 45660 characters]\" align=\"left\" width=\"400\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca411d89",
   "metadata": {},
   "source": [
    "For the Laplace equation in 2D, the interior values in the computational grid (represented by a matrix $u$) are computed with this iterative scheme. The entry $(i,j)$ of  matrix $u$ at iteration $t+1$ is computed as:\n",
    "\n",
    "\n",
    "$u^{t+1}_{(i,j)} = \\dfrac{u^t_{(i-1,j)}+u^t_{(i+1,j)}+u^t_{(i,j-1)}+u^t_{(i,j+1)}}{4}$\n",
    "\n",
    "Note that each entry is updated as the average of the four neighbors (top,bottom,left,right) of that entry in the previous iteration. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c916627",
   "metadata": {},
   "source": [
    "### Serial implementation\n",
    "\n",
    "The next code implements a simple example, where the boundary values are equal to 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e843565",
   "metadata": {},
   "outputs": [],
   "source": [
    "function jacobi_2d(n,niters)\n",
    "    u = zeros(n+2,n+2)\n",
    "    u[1,:] = u[end,:] = u[:,1] = u[:,end] .= 1\n",
    "    u_new = copy(u)\n",
    "    for t in 1:niters\n",
    "        for j in 2:(n+1)\n",
    "            for i in 2:(n+1)\n",
    "                north = u[i,j+1]\n",
    "                south = u[i,j-1]\n",
    "                east = u[i+1,j]\n",
    "                west = u[i-1,j]\n",
    "                u_new[i,j] = 0.25*(north+south+east+west)\n",
    "            end\n",
    "        end\n",
    "        u, u_new = u_new, u\n",
    "    end\n",
    "    u\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4dc50b1",
   "metadata": {},
   "source": [
    "If you run the function for zero iterations you will see the initial condition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4acf219e",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10\n",
    "niters = 0\n",
    "u = jacobi_2d(n,niters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efd4df6a",
   "metadata": {},
   "source": [
    "If you run the problem for enough iterations, you should converge to the exact solution. In this case, the exact solution is a matrix with all values equal to one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c120b8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10\n",
    "niters = 300\n",
    "u = jacobi_2d(n,niters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85895681",
   "metadata": {},
   "source": [
    "### Where can we exploit parallelism?\n",
    "\n",
    "```julia\n",
    "for t in 1:niters\n",
    "    for j in 2:(n+1)\n",
    "        for i in 2:(n+1)\n",
    "            north = u[i,j+1]\n",
    "            south = u[i,j-1]\n",
    "            east = u[i+1,j]\n",
    "            west = u[i-1,j]\n",
    "            u_new[i,j] = 0.25*(north+south+east+west)\n",
    "        end\n",
    "    end\n",
    "    u, u_new = u_new, u\n",
    "end\n",
    "```\n",
    "\n",
    "- The outer loop cannot be parallelized (like in the 1d case). \n",
    "- The two inner loops are trivially parallel\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65e4f745",
   "metadata": {},
   "source": [
    "### Parallelization strategies\n",
    "\n",
    "In 2d one has more flexibility in order to distribute the data over the processes. We consider these three alternatives:\n",
    "\n",
    "- 1D block row partition (each worker handles a subset of consecutive rows and all columns)\n",
    "- 2D block partition (each worker handles a subset of consecutive rows and columns)\n",
    "- 2D cyclic partition (each workers handles a subset of alternating rows ans columns)\n",
    "\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b> Other options are 1D block column partition and 1D cyclic (row or column) partition. They are not analyzed in this notebook since they are closely related to the other strategies. In Julia, in fact, it is often preferable to work with 1D block column partitions than with 1D block row partitions since matrices are stored in column major order.\n",
    "</div>\n",
    "\n",
    "\n",
    "The three partition types are depicted in the following figure for 4 processes."
   ]
  },
  {
   "attachments": {
    "fig-jacobi-partitions.png": {
     "image/png": "[encoded data, 30376 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "11e834b5",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 30376 characters]\" align=\"left\" width=\"450\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7384ddfe",
   "metadata": {},
   "source": [
    "Which of the thee alternatives is more efficient? To answer this question we need to quantify how much data is processed and communicated in each case. The following analysis assumes that the grid is of $N$ by $N$ cells and that the number of processes is $P$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f733d88f",
   "metadata": {},
   "source": [
    "### 1D block partition\n",
    "\n",
    "The following figure shows the portion of vector `u_new` that is updated at each iteration by a particular process (CPU 3) left picture, and which entries of `u` are needed to update this data, right picture. We use analogous figures for the other partitions below.\n"
   ]
  },
  {
   "attachments": {
    "g26521.png": {
     "image/png": "[encoded data, 20120 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "6b983ffd",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 20120 characters]\" align=\"left\" width=\"400\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bc21623",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "<b>Question:</b>  Compute the complexity of the communication over computation ratio for this data partition.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d01f8ce8",
   "metadata": {},
   "outputs": [],
   "source": [
    "uncover = false # Change to true to see the answer\n",
    "partition_1d_answer(uncover)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9bed431",
   "metadata": {},
   "source": [
    "### 2D block partition"
   ]
  },
  {
   "attachments": {
    "g26305.png": {
     "image/png": "[encoded data, 19492 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "6f7b5b36",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 19492 characters]\" align=\"left\" width=\"400\">\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09bd28ca",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "<b>Question:</b>  Compute the complexity of the communication over computation ratio for this data partition.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e94a1ea6",
   "metadata": {},
   "outputs": [],
   "source": [
    "uncover = false\n",
    "partition_2d_answer(uncover)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9aaf1ca6",
   "metadata": {},
   "source": [
    "### 2D cyclic partition"
   ]
  },
  {
   "attachments": {
    "g26911.png": {
     "image/png": "[encoded data, 22796 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "77701278",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 22796 characters]\" align=\"left\" width=\"400\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b373e9ce",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "<b>Question:</b>  Compute the complexity of the communication over computation ratio for this data partition.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10fab825",
   "metadata": {},
   "outputs": [],
   "source": [
    "uncover = false\n",
    "partition_cyclic_answer(uncover)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "191fa03b",
   "metadata": {},
   "source": [
    "### Summary\n",
    "\n",
    "|Partition | Messages <br> per iteration | Communication <br>per worker | Computation <br>per worker | Ratio communication/<br>computation |\n",
    "|---|---|---|---|---|\n",
    "| 1d block | 2 | O(N) | N²/P | O(P/N) |\n",
    "| 2d block | 4 | O(N/√P) | N²/P | O(√P/N) |\n",
    "| 2d cyclic | 4 |O(N²/P) | N²/P | O(1) |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0c6d216",
   "metadata": {},
   "source": [
    "### Which partition is the best one?\n",
    "\n",
    "\n",
    "\n",
    "- Both 1d and 2d block partitions are potentially scalable if $P<<N$\n",
    "- The 2d block partition has the lowest communication complexity\n",
    "- The 1d block partition requires to send less messages (It can be useful if the fixed cost of sending a message is high)\n",
    "- The best strategy for a given problem size will thus depend on the machine.\n",
    "- Cyclic partitions are impractical for this application (but they are useful in others) \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8caf9b30",
   "metadata": {},
   "source": [
    "## The Gauss-Seidel method \n",
    "\n",
    "Now let us study a slightly more challenging method. The implementation of Jacobi used an auxiliary array `u_new`. The usage of `u_new` seems a bit unnecessary at first sight, right? If we remove it, we get another method called Gauss-Seidel.\n",
    "\n",
    "\n",
    "The following cell contains the sequential implementation of Gauss-Seidel. It is obtained taking the sequential implementation of Jacobi, removing `u_new`, and only using `u`. This method is more efficient in terms of memory consumed (only one vector required instead of two). However it is harder to parallelize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efd792c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "function gauss_seidel(n,niters)\n",
    "    u = zeros(n+2)\n",
    "    u[1] = -1\n",
    "    u[end] = 1\n",
    "    for t in 1:niters\n",
    "        for i in 2:(n+1)\n",
    "            u[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "    end\n",
    "    u\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a38cef6",
   "metadata": {},
   "source": [
    "Note that the final solution is nearly the same as for Jacobi for a large enough number of iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7efa254",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5\n",
    "niters = 1000\n",
    "gauss_seidel(n,niters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37b7c288",
   "metadata": {},
   "outputs": [],
   "source": [
    "jacobi(n,niters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97d9ad59",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "<b>Question:</b>  Which of the two loops in the Gauss-Seidel method are trivially parallelizable?\n",
    "</div>\n",
    "\n",
    "```julia\n",
    "for t in 1:niters\n",
    "    for i in 2:(n+1)\n",
    "        u[i] = 0.5*(u[i-1]+u[i+1])\n",
    "    end\n",
    "end\n",
    "```\n",
    "\n",
    "    a) Both of them\n",
    "    b) The outer, but not the inner\n",
    "    c) None of them\n",
    "    d) The inner, but not the outer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b383598",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "answer = \"x\" # replace x with a, b, c or d\n",
    "gauss_seidel_1_check(answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b06f709",
   "metadata": {},
   "source": [
    "If you look into the algorithm closely, you will realize that the steps in loop over `i` are not independent. To update `u[i]` we first need to update `u[i-1]`, which requires to update `u[i-2]` etc. This happens for Gauss-Seidel but not for Jacobi. For Jacobi the updates are written into the temporary array `u_new` instead of updating `u` directly. In other words, in Jacobi, updating `u_new[i]` does not require any entry updated entry of `u_new`. It only requires entries in vector `u`, which is not modified in the loop over `i`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "099634d9",
   "metadata": {},
   "source": [
    "### Backwards Gauss-Seidel\n",
    "\n",
    "In addition, the the result of the Gauss-Seidel method depends on the order of the steps in the loop over `i`. This is another symptom that tells you that this loop is hard (or impossible) to parallelize. For instance, if you do the iterations over `i` by reversing the loop order, you get another method called *backward* Gauss-Seidel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff03c246",
   "metadata": {},
   "outputs": [],
   "source": [
    "function backward_gauss_seidel(n,niters)\n",
    "    u = zeros(n+2)\n",
    "    u[1] = -1\n",
    "    u[end] = 1\n",
    "    for t in 1:niters\n",
    "        for i in (n+1):-1:2\n",
    "            u[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "    end\n",
    "    u\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63c4ce1f",
   "metadata": {},
   "source": [
    "Both Jacobi and *forward* and *backward* Gauss-Seidel converge to the same result, but they lead to slightly different values during the iterations. Check it with the following cells. First, run the methods with `niters=1` and then with `niters=100`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abe31423",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5\n",
    "niters = 1\n",
    "u_forward = gauss_seidel(n,niters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23430f70",
   "metadata": {},
   "outputs": [],
   "source": [
    "u_backward = backward_gauss_seidel(n,niters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b2b2437",
   "metadata": {},
   "outputs": [],
   "source": [
    "u_jacobi = jacobi(n,niters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0deb5549",
   "metadata": {},
   "source": [
    "### Red-black Gauss-Seidel\n",
    "\n",
    "There is yet another version called *red-black* Gauss-Seidel. This uses a very clever order for the steps in the loop over `i`. It does this loop in two phases. First, one updates the entries with even index, and then the entries with odd index."
   ]
  },
  {
   "attachments": {
    "g28304.png": {
     "image/png": "[encoded data, 28024 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "dda0a371",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 28024 characters]\"  align=\"left\" width=\"300\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9707f13b",
   "metadata": {},
   "source": [
    "The implementation is given in the following cell. Note that we introduced an extra loop that represent\n",
    "each one of the two phases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0cbe2d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "function red_black_gauss_seidel(n,niters)\n",
    "    u = zeros(n+2)\n",
    "    u[1] = -1\n",
    "    u[end] = 1\n",
    "    for t in 1:niters\n",
    "        for color in (0,1)\n",
    "            for i in (n+1):-1:2\n",
    "                if color == mod(i,2)\n",
    "                    u[i] = 0.5*(u[i-1]+u[i+1])\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    u\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13ce270f",
   "metadata": {},
   "source": [
    "Run the method for several values of `niters`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59c1e7f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5\n",
    "niters = 1000\n",
    "red_black_gauss_seidel(n,niters)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1bacc26",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "<b>Question:</b>  Which of the loops in the red-black Gauss-Seidel method are trivially parallelizable?\n",
    "</div>\n",
    "\n",
    "```julia\n",
    "for t in 1:niters\n",
    "    for color in (0,1)\n",
    "        for i in (n+1):-1:2\n",
    "            if color == mod(i,2)\n",
    "                u[i] = 0.5*(u[i-1]+u[i+1])\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "end\n",
    "```\n",
    "\n",
    "    a) All loops\n",
    "    b) Loop over t only\n",
    "    c) Loop over color only\n",
    "    d) Loop over i only\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d511cc02",
   "metadata": {},
   "outputs": [],
   "source": [
    "answer = \"x\" # replace x with a, b, c or d\n",
    "gauss_seidel_2_check(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8946b83f",
   "metadata": {},
   "outputs": [],
   "source": [
    "uncover = false\n",
    "gauss_seidel_2_why(uncover)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41e90d60",
   "metadata": {},
   "source": [
    "### Changing an algorithm to make it parallel\n",
    "\n",
    "Note that the original method (the forward Gauss-Seidel) cannot be parallelized, we needed to modify the method slightly with the red-black ordering in order to create a method that can be parallelized. However the method we parallelized is not equivalent to the original one. This happens in practice in many other applications. An algorithm might be impossible to parallelize and one needs to modify it to exploit parallelism. However, one needs to be careful when modifying the algorithm to not destroy the algorithmic properties of the original one. In this case, we succeeded. The red-black Gauss-Seidel converges as fast (if not faster) than the original forward Gauss-Seidel. However, this is not true in general. There is often a trade-off between the algorithmic properties and how parallelizable is the algorithm."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ed4129c",
   "metadata": {},
   "source": [
    "## MPI implementation\n",
    "\n",
    "In the last part of this notebook, we consider the implementation of the Jacobi method using MPI. We will consider the 1D version for simplicity.\n",
    "\n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b> The programming model of MPI is generally better suited for data-parallel algorithms like this one than the task-based model provided by Distributed.jl. In any case, one can also implement it using Distributed.jl, but it requires some extra effort to setup the remote channels right for the communication between neighbor processes. \n",
    "</div>\n",
    "\n",
    "\n",
    "We are going to implement the method in a function with the following signature\n",
    "\n",
    "```julia\n",
    "    u = jacobi_mpi(n,niters,comm)\n",
    "```\n",
    "\n",
    "The signature will be very similar to the sequential function `jacobi`, but there are some important differences:\n",
    "\n",
    "1. The parallel one takes an MPI communicator object. This allows the user to decide which communicator to use. For instance, `MPI.COMM_WORLD` directly or a duplicate of this one (which is the recommended approach).\n",
    "2. Function `jacobi_mpi` will be called on multiple MPI ranks. Thus, the values of `n` and `niters` should be the same in all ranks. The caller of this function has to make sure that this is true. Otherwise, the parallel code will not work.\n",
    "3. The result `u` is NOT the same as the result of the sequential function `jacobi`. It will contained only the local portion of `u` stored at each rank. To recover the same vector as in the sequential case, one needs to gather these pieces in a single rank. We will do this later.\n",
    "\n",
    "The implementation of the parallel function is given in the following code snipped:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d31c1e41",
   "metadata": {},
   "outputs": [],
   "source": [
    "code1 = quote\n",
    "    function jacobi_mpi(n,niters,comm)\n",
    "        # Initialization\n",
    "        u, u_new = init(n,comm)\n",
    "        for t in 1:niters\n",
    "            # Communication\n",
    "            ghost_exchange!(u,comm)\n",
    "            # Local computation\n",
    "            local_update!(u,u_new)\n",
    "            u, u_new = u_new, u\n",
    "        end\n",
    "        return u\n",
    "    end\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04a38c37",
   "metadata": {},
   "source": [
    "In the following cells, we discuss the implementation of the helper functions `init`, `ghost_exchange`, and `local_update`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96b8d40e",
   "metadata": {},
   "source": [
    "### Initialization\n",
    "\n",
    "Let us start with function `init`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b98a9348",
   "metadata": {},
   "outputs": [],
   "source": [
    "code2 = quote\n",
    "    function init(n,comm)\n",
    "        nranks = MPI.Comm_size(comm)\n",
    "        rank = MPI.Comm_rank(comm)\n",
    "        if mod(n,nranks) != 0\n",
    "            println(\"n must be a multiple of nranks\")\n",
    "            MPI.Abort(comm,1)\n",
    "        end\n",
    "        load = div(n,nranks)\n",
    "        u = zeros(load+2)\n",
    "        if rank == 0\n",
    "            u[1] = -1\n",
    "        end\n",
    "        if rank == nranks-1\n",
    "            u[end] = 1\n",
    "        end\n",
    "        u_new = copy(u)\n",
    "        u, u_new\n",
    "    end\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b9e75d8",
   "metadata": {},
   "source": [
    "This function crates and initializes the vector `u` and the auxiliary vector `u_new` and fills in the boundary values. Note that we are not creating the full arrays like in the sequential case. We are only creating the parts to be managed by the current rank. To this end, we start by computing the number of entries to be updated in this rank, i.e., variable `load`. We have assumed that `n` is a multiple of the number of ranks for simplicity. If this is not the case, we stop the computation with `MPI.Abort`. Note that we are allocating two extra elements in `u` (and `u_new`) for the ghost cells and boundary conditions. The following figure displays the arrays created for `n==9` and `nranks==3` (thus `load == 3`). Note that the first and last elements of the arrays are displayed with gray edges denoting that they are the extra elements allocated for ghost cells or boundary conditions."
   ]
  },
  {
   "attachments": {
    "fig.png": {
     "image/png": "[encoded data, 21324 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "756b8d01",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 21324 characters]\"  align=\"left\" width=\"450\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f0e49e4",
   "metadata": {},
   "source": [
    "### Communication\n",
    "\n",
    "Once the working arrays have been created, we can start with the iterations of the Jacobi method. Remember that this is implemented in parallel by updating the ghost cells with data from neighboring processes and then performing the sequential Jacobi update. See figure:"
   ]
  },
  {
   "attachments": {
    "fig15.png": {
     "image/png": "[encoded data, 29440 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "6b214777",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 29440 characters]\"  align=\"left\" width=\"450\"/>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1bd8413",
   "metadata": {},
   "source": [
    "The communication step happens in function `ghost_exchange!`. This one modifies the ghost cells in the input vector `u` by importing data using MPI point-to-point communication. See the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "688638a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "code3 = quote\n",
    "    function ghost_exchange!(u,comm)\n",
    "        load = length(u)-2\n",
    "        rank = MPI.Comm_rank(comm)\n",
    "        nranks = MPI.Comm_size(comm)\n",
    "        if rank != 0\n",
    "            neig_rank = rank-1\n",
    "            u_snd = view(u,2:2)\n",
    "            u_rcv = view(u,1:1)\n",
    "            dest = neig_rank\n",
    "            source = neig_rank\n",
    "            MPI.Sendrecv!(u_snd,u_rcv,comm;dest,source)\n",
    "        end\n",
    "        if rank != (nranks-1)\n",
    "            neig_rank = rank+1\n",
    "            u_snd = view(u,(load+1):(load+1))\n",
    "            u_rcv = view(u,(load+2):(load+2))\n",
    "            dest = neig_rank\n",
    "            source = neig_rank\n",
    "            MPI.Sendrecv!(u_snd,u_rcv,comm;dest,source)\n",
    "        end\n",
    "    end\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bf5408c",
   "metadata": {},
   "source": [
    "Note that we have used `MPI.Sendrecv!` to send and receive values."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2ce54e4",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "<b>Question:</b>  Is this other implementation based on `MPI.Send` and `MPI.Recv!` correct?\n",
    "</div>\n",
    "\n",
    "```julia\n",
    "    function ghost_exchange!(u,comm)\n",
    "        load = length(u)-2\n",
    "        rank = MPI.Comm_rank(comm)\n",
    "        nranks = MPI.Comm_size(comm)\n",
    "        if rank != 0\n",
    "            neig_rank = rank-1\n",
    "            u_snd = view(u,2:2)\n",
    "            u_rcv = view(u,1:1)\n",
    "            dest = neig_rank\n",
    "            source = neig_rank\n",
    "            MPI.Send(u_snd,comm;dest)\n",
    "            MPI.Recv!(u_rcv,comm;source)\n",
    "        end\n",
    "        if rank != (nranks-1)\n",
    "            neig_rank = rank+1\n",
    "            u_snd = view(u,(load+1):(load+1))\n",
    "            u_rcv = view(u,(load+2):(load+2))\n",
    "            dest = neig_rank\n",
    "            source = neig_rank\n",
    "            MPI.Send(u_snd,comm;dest)\n",
    "            MPI.Recv!(u_rcv,comm;source)\n",
    "        end\n",
    "    end\n",
    "```\n",
    "\n",
    "    a) It is correct.\n",
    "    c) It is incorrect, but it might provide the right result depending on the MPI implementation.\n",
    "    b) It is incorrect, and it is guaranteed that it will result in a dead lock.\n",
    "    d) This implementation does not work when distributing over just a single MPI rank.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8498c7e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "answer = \"x\" # replace x with a, b, c or d\n",
    "sndrcv_check(answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d863cb89",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "<b>Question:</b>  How would you fix the implementation above, while still using `MPI.Send` and `MPI.Recv!` instead of `MPI.Sendrecv!` ?\n",
    "</div>\n",
    "\n",
    "Think about the answer. To uncover the the solution run next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "629a3356",
   "metadata": {},
   "outputs": [],
   "source": [
    "uncover = false\n",
    "sndrcv_fix_answer(uncover)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82e21f4e",
   "metadata": {},
   "source": [
    "### Local computation\n",
    "\n",
    "Once the ghost cells have the right values, we can perform the Jacobi update locally at each process. This is done in function `local_update!`. Note that here we only update the data *owned* by the current MPI rank, i.e. we do not modify the ghost values. There is no need to modify the ghost values since they will updated by another rank. In the code, this is reflected in the loop over `i`. We do not visit the first nor the last entry in `u_new`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f43560e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "code4 = quote\n",
    "    function local_update!(u,u_new)\n",
    "        load = length(u)-2\n",
    "        for i in 2:(load+1)\n",
    "            u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "    end\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47861c19",
   "metadata": {},
   "source": [
    "### Running the code\n",
    "\n",
    "Let us put all pieces together and run the code. If not done yet, install MPI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53674411",
   "metadata": {},
   "outputs": [],
   "source": [
    "] add MPI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c966375a",
   "metadata": {},
   "source": [
    "The following cells includes all previous code snippets into a final one. We are eventually calling function `jacobi_mpi` and showing the result vector `u`. Run the following code for 1 MPI rank, then for 2 and 3 MPI ranks. Look into the values of `u`. Does it make sense?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f91e4759",
   "metadata": {},
   "outputs": [],
   "source": [
    "using MPI\n",
    "code = quote\n",
    "    using MPI\n",
    "    MPI.Init()\n",
    "    $code1\n",
    "    $code2\n",
    "    $code3\n",
    "    $code4\n",
    "    n = 9\n",
    "    niters = 200\n",
    "    comm = MPI.Comm_dup(MPI.COMM_WORLD)\n",
    "    u = jacobi_mpi(n,niters,comm)\n",
    "    @show u\n",
    "end\n",
    "run(`$(mpiexec()) -np 1 julia --project=. -e $code`);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e4fe21d",
   "metadata": {},
   "source": [
    "### Checking the result\n",
    "\n",
    "Checking the result visually is not enough in general. To check the parallel implementation we want to compare it against the sequential implementation. However, how can we compare the sequential and the parallel result? The parallel version gives a distributed vector. We cannot compare this one directly with the result of the sequential function. A possible solution is to gather all the pieces of the parallel result in a single rank and compare there against the parallel implementation.\n",
    "\n",
    "\n",
    "The following function gather the distributed vector in rank 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49a220c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "code5 = quote\n",
    "    function gather_final_result(u,comm)\n",
    "        load = length(u)-2\n",
    "        rank = MPI.Comm_rank(comm)\n",
    "        if rank !=0\n",
    "            # If I am not rank 0\n",
    "            # I send my own data to rank 0\n",
    "            u_snd = view(u,2:(load+1))\n",
    "            MPI.Send(u_snd,comm,dest=0)\n",
    "            u_root = zeros(0) # This will nevel be used\n",
    "        else\n",
    "            # If I am rank 0\n",
    "            nranks = MPI.Comm_size(comm)\n",
    "            n = load*nranks\n",
    "            u_root = zeros(n+2)\n",
    "            # Set boundary\n",
    "            u_root[1] = -1\n",
    "            u_root[end] = 1\n",
    "            # Set data for rank 0\n",
    "            lb = 2\n",
    "            ub = load+1\n",
    "            u_root[lb:ub] = view(u,lb:ub)\n",
    "            # Receive and set data from other ranks\n",
    "            for other_rank in 1:(nranks-1)\n",
    "                lb += load\n",
    "                ub += load\n",
    "                u_rcv = view(u_root,lb:ub)\n",
    "                MPI.Recv!(u_rcv,comm;source=other_rank)\n",
    "            end\n",
    "        end\n",
    "        # NB onle the root (rank 0) will\n",
    "        # contain meaningfull data\n",
    "        return u_root\n",
    "    end\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e66bad1b",
   "metadata": {},
   "source": [
    "Run the following cell to see the result. Is the result as expected?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a9f4a54",
   "metadata": {},
   "outputs": [],
   "source": [
    "code = quote\n",
    "    using MPI\n",
    "    MPI.Init()\n",
    "    $code1\n",
    "    $code2\n",
    "    $code3\n",
    "    $code4\n",
    "    $code5\n",
    "    n = 9\n",
    "    niters = 200\n",
    "    comm = MPI.Comm_dup(MPI.COMM_WORLD)\n",
    "    u = jacobi_mpi(n,niters,comm)\n",
    "    u_root = gather_final_result(u,comm)\n",
    "    @show u_root\n",
    "end\n",
    "run(`$(mpiexec()) -np 3 julia --project=. -e $code`);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc9b2b51",
   "metadata": {},
   "source": [
    "Now that we have collected the parallel vector into a single array, we can compare it against the sequential implementation. This is done in the following cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb053382",
   "metadata": {},
   "outputs": [],
   "source": [
    "code6 = quote\n",
    "    function jacobi(n,niters)\n",
    "        u = zeros(n+2)\n",
    "        u[1] = -1\n",
    "        u[end] = 1\n",
    "        u_new = copy(u)\n",
    "        for t in 1:niters\n",
    "            for i in 2:(n+1)\n",
    "                u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "            end\n",
    "            u, u_new = u_new, u\n",
    "        end\n",
    "        u\n",
    "    end\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffe6c605",
   "metadata": {},
   "outputs": [],
   "source": [
    "code = quote\n",
    "    using MPI\n",
    "    MPI.Init()\n",
    "    $code1\n",
    "    $code2\n",
    "    $code3\n",
    "    $code4\n",
    "    $code5\n",
    "    $code6\n",
    "    n = 12\n",
    "    niters = 100\n",
    "    comm = MPI.Comm_dup(MPI.COMM_WORLD)\n",
    "    u = jacobi_mpi(n,niters,comm)\n",
    "    u_root = gather_final_result(u,comm)\n",
    "    rank = MPI.Comm_rank(comm)\n",
    "    if rank == 0\n",
    "        # Compare agains serial\n",
    "        u_seq = jacobi(n,niters)\n",
    "        if isapprox(u_root,u_seq)\n",
    "            println(\"Test passed 🥳\")\n",
    "        else\n",
    "            println(\"Test failed\")\n",
    "        end\n",
    "    end\n",
    "end\n",
    "run(`$(mpiexec()) -np 3 julia --project=. -e $code`);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9aa2901",
   "metadata": {},
   "source": [
    "## Latency hiding\n",
    "\n",
    "We have now a correct parallel implementation. But. can our implementation above be improved? Note that we only need communications to update the values at the boundary of the portion owned by each process. The other values (the one in green in the figure below) can be updated without communications. This provides the opportunity of overlapping the computation of the interior values (green cells in the figure) with the communication of the ghost values. This technique is called latency hiding, since we are hiding communication latency by overlapping it with computation that we need to do anyway. The actual implementation is left as an exercise (see Exercise 1)."
   ]
  },
  {
   "attachments": {
    "fig16.png": {
     "image/png": "[encoded data, 12140 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "7d66b1a2",
   "metadata": {},
   "source": [
    "<div>\n",
    "<img src=\"data:image/png;base64,[encoded data, 12140 characters]\" align=\"left\" width=\"450\"/>\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2d48c45",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">\n",
    "<b>Question:</b>  Which MPI directives would you use to implement latency hiding in the communication the ghost values?\n",
    "</div>\n",
    "\n",
    "    a) MPI_Send and MPI_Recv\n",
    "    b) MPI_Bsend and MPI_Recv\n",
    "    c) MPI_Isend and MPI_Irecv \n",
    "    d) MPI_Sendrecv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29ade8f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "answer = \"x\" # replace x with a, b, c or d\n",
    "lh_check(answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bd014b0",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "- We learned how to parallelize the Jacobi method efficiently.\n",
    "- Using block partitions the communication overhead is small if the problem size is larger than the number of processors.\n",
    "- In addition, one can use latency hiding to further reduce the overhead caused by communication.\n",
    "- Both are not true if we use cyclic data partitions.\n",
    "- Gauss-Seidel is a simple variation of Jacobi but is much more challenging to parallelize.\n",
    "- One needs to consider a special order (red-black) when updating the values to parallelize the method.\n",
    "- We learned how to implement a 1D Jacobi method with MPI.\n",
    "- Once needs to be careful when using blocking directives to avoid dead locks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47643bf6",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6527d24f",
   "metadata": {},
   "source": [
    "### Exercise 1\n",
    "\n",
    "The following code implements the 1D Jacobi method studied in this notebook, but using non-blocking sends and receives. Modify this code so that you overlap the communication of the ghost values with local computations as explained above in the section \"latency hiding\". You only need to modify function `jacobi_mpi_latency_hiding`. Copy the code below to a file called `ex1.jl`. Modify the file (e.g. with vscode).  Run it from the Julia REPL using the `run` function as explained in the [Getting Started tutorial](https://www.francescverdugo.com/XM_40017/dev/getting_started_with_julia/#Running-MPI-code)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b57f876",
   "metadata": {},
   "source": [
    "\n",
    "```julia\n",
    "# file ex1.jl (begin)\n",
    "\n",
    "using MPI\n",
    "MPI.Init()\n",
    "\n",
    "function jacobi_mpi_latency_hiding(n,niters,comm)\n",
    "    u, u_new = init(n,comm)\n",
    "    load = length(u)-2\n",
    "    rank = MPI.Comm_rank(comm)\n",
    "    nranks = MPI.Comm_size(comm)\n",
    "    nreqs = 2*((rank != 0) + (rank != (nranks-1)))\n",
    "    reqs = MPI.MultiRequest(nreqs)\n",
    "    for t in 1:niters\n",
    "        ireq = 0\n",
    "        if rank != 0\n",
    "            neig_rank = rank-1\n",
    "            u_snd = view(u,2:2)\n",
    "            u_rcv = view(u,1:1)\n",
    "            dest = neig_rank\n",
    "            source = neig_rank\n",
    "            ireq += 1\n",
    "            MPI.Isend(u_snd,comm,reqs[ireq];dest)\n",
    "            ireq += 1\n",
    "            MPI.Irecv!(u_rcv,comm,reqs[ireq];source)\n",
    "        end\n",
    "        if rank != (nranks-1)\n",
    "            neig_rank = rank+1\n",
    "            u_snd = view(u,(load+1):(load+1))\n",
    "            u_rcv = view(u,(load+2):(load+2))\n",
    "            dest = neig_rank\n",
    "            source = neig_rank\n",
    "            ireq += 1\n",
    "            MPI.Isend(u_snd,comm,reqs[ireq];dest)\n",
    "            ireq += 1\n",
    "            MPI.Irecv!(u_rcv,comm,reqs[ireq];source)\n",
    "        end\n",
    "        MPI.Waitall(reqs)\n",
    "        for i in 2:load+1\n",
    "            u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "        u, u_new = u_new, u\n",
    "    end\n",
    "    return u\n",
    "end\n",
    "\n",
    "function init(n,comm)\n",
    "    nranks = MPI.Comm_size(comm)\n",
    "    rank = MPI.Comm_rank(comm)\n",
    "    if mod(n,nranks) != 0\n",
    "        println(\"n must be a multiple of nranks\")\n",
    "        MPI.Abort(comm,1)\n",
    "    end\n",
    "    load = div(n,nranks)\n",
    "    u = zeros(load+2)\n",
    "    if rank == 0\n",
    "        u[1] = -1\n",
    "    end\n",
    "    if rank == nranks-1\n",
    "        u[end] = 1\n",
    "    end\n",
    "    u_new = copy(u)\n",
    "    u, u_new\n",
    "end\n",
    "\n",
    "function gather_final_result(u,comm)\n",
    "    load = length(u)-2\n",
    "    rank = MPI.Comm_rank(comm)\n",
    "    if rank !=0\n",
    "        u_snd = view(u,2:(load+1))\n",
    "        MPI.Send(u_snd,comm,dest=0)\n",
    "        u_root = zeros(0)\n",
    "    else\n",
    "        nranks = MPI.Comm_size(comm)\n",
    "        n = load*nranks\n",
    "        u_root = zeros(n+2)\n",
    "        u_root[1] = -1\n",
    "        u_root[end] = 1\n",
    "        lb = 2\n",
    "        ub = load+1\n",
    "        u_root[lb:ub] = view(u,lb:ub)\n",
    "        for other_rank in 1:(nranks-1)\n",
    "            lb += load\n",
    "            ub += load\n",
    "            u_rcv = view(u_root,lb:ub)\n",
    "            MPI.Recv!(u_rcv,comm;source=other_rank)\n",
    "        end\n",
    "    end\n",
    "    return u_root\n",
    "end\n",
    "\n",
    "function jacobi(n,niters)\n",
    "    u = zeros(n+2)\n",
    "    u[1] = -1\n",
    "    u[end] = 1\n",
    "    u_new = copy(u)\n",
    "    for t in 1:niters\n",
    "        for i in 2:(n+1)\n",
    "            u_new[i] = 0.5*(u[i-1]+u[i+1])\n",
    "        end\n",
    "        u, u_new = u_new, u\n",
    "    end\n",
    "    u\n",
    "end\n",
    "\n",
    "n = 12\n",
    "niters = 100\n",
    "comm = MPI.Comm_dup(MPI.COMM_WORLD)\n",
    "u = jacobi_mpi_latency_hiding(n,niters,comm)\n",
    "u_root = gather_final_result(u,comm)\n",
    "rank = MPI.Comm_rank(comm)\n",
    "if rank == 0\n",
    "    u_seq = jacobi(n,niters)\n",
    "    if isapprox(u_root,u_seq)\n",
    "        println(\"Test passed 🥳\")\n",
    "    else\n",
    "        println(\"Test failed 😢\")\n",
    "    end\n",
    "end\n",
    "\n",
    "# file ex1.jl (end)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1af362aa",
   "metadata": {},
   "source": [
    "### Exercise 2\n",
    "\n",
    "In the parallel implementation of the Jacobi method, we assumed that the method runs for a given number of iterations. In function, `jacobi_with_tol` at the beginning of the notebook shows how the Jacobi iterations can be stopped when the difference between iterations is small. Implement a parallel version of this function. Start with the in Exercise 1 and add the stopping criterion implemented in `jacobi_with_tol`. Use a text editor and the Julia REPL. Do not try to implement the code in a notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d3430ad",
   "metadata": {},
   "source": [
    "# License\n",
    "\n",
    "\n",
    "\n",
    "This notebook is part of the course [Programming Large Scale Parallel Systems](https://www.francescverdugo.com/XM_40017) at Vrije Universiteit Amsterdam and may be used under a [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/) license."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.10.0",
   "language": "julia",
   "name": "julia-1.10"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
