diff --git a/notebooks/8__mpi.ipynb b/notebooks/8__mpi.ipynb new file mode 100644 index 0000000..932f9d3 --- /dev/null +++ b/notebooks/8__mpi.ipynb @@ -0,0 +1,570 @@ +{ + "cells": [ + { + "cell_type": "code", + "source": [ + "## Uncomment and Execute this cell if running on slurm based cluster\n", + "#SRUN=\"srun\"" + ], + "metadata": { + "id": "P8UcZACXA4i4" + }, + "id": "P8UcZACXA4i4", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "## Uncomment and Execute this cell if running on google collab\n", + "#!pip install mpi4py\n", + "#SRUN=\"mpirun --allow-run-as-root --oversubscribe\"" + ], + "metadata": { + "id": "0sW2XPMR-687" + }, + "id": "0sW2XPMR-687", + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "32b1f077-779c-432d-b807-98f5646fd9f4", + "metadata": { + "id": "32b1f077-779c-432d-b807-98f5646fd9f4" + }, + "outputs": [], + "source": [ + "%%writefile hello_world.py\n", + "\n", + "from mpi4py import MPI\n", + "\n", + "print(\"Hello World!\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11291ac1-088b-4c97-b32c-9dee22f0e1d1", + "metadata": { + "id": "11291ac1-088b-4c97-b32c-9dee22f0e1d1" + }, + "outputs": [], + "source": [ + "!$SRUN --allow-run-as-root --oversubscribe -n 2 -c 2 python hello_world.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ca7a23c9-9bbf-489b-b697-a10b9e24ebc0", + "metadata": { + "id": "ca7a23c9-9bbf-489b-b697-a10b9e24ebc0" + }, + "outputs": [], + "source": [ + "!$SRUN -n 4 -c 4 python hello_world.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "06e5d0d5-797f-40b1-a69d-365126f2ea67", + "metadata": { + "id": "06e5d0d5-797f-40b1-a69d-365126f2ea67" + }, + "outputs": [], + "source": [ + "%%writefile hello_world_with_ranks.py\n", + "import os\n", + "from mpi4py import MPI\n", + "comm = MPI.COMM_WORLD\n", + "size = comm.Get_size()\n", + "rank = comm.Get_rank()\n", + "name = MPI.Get_processor_name()\n", + "pid = os.getpid()\n", + "print(f\"Hello World! I am process {rank} of {size} on {name} \"\n", + " \"with pid {pid}.\\n\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bfa7013c-c3ee-4a74-a7e3-94985ce04c0b", + "metadata": { + "id": "bfa7013c-c3ee-4a74-a7e3-94985ce04c0b" + }, + "outputs": [], + "source": [ + "!$SRUN -n 2 -c 2 python hello_world_with_ranks.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1a032e45-682a-4c55-9a34-e7cdc01f8a2d", + "metadata": { + "id": "1a032e45-682a-4c55-9a34-e7cdc01f8a2d" + }, + "outputs": [], + "source": [ + "%%writefile send_recv.py\n", + "\n", + "from mpi4py import MPI\n", + "\n", + "comm = MPI.COMM_WORLD\n", + "size = comm.Get_size()\n", + "rank = comm.Get_rank()\n", + "if rank == 0:\n", + " data = {'a': 7, 'b': 3.14}\n", + " comm.send(data, dest=1, tag=11)\n", + "elif rank == 1:\n", + " data = comm.recv(source=0, tag=11)\n", + " print(data)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ab2db951-10f5-4570-8717-d9c4ed88acab", + "metadata": { + "id": "ab2db951-10f5-4570-8717-d9c4ed88acab" + }, + "outputs": [], + "source": [ + "!$SRUN -n 2 -c 2 python send_recv.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ab6096d9-d134-490e-9c92-3f800d9a5c8e", + "metadata": { + "id": "ab6096d9-d134-490e-9c92-3f800d9a5c8e" + }, + "outputs": [], + "source": [ + "%%writefile send_recv_eff.py\n", + "import numpy as np\n", + "from mpi4py import MPI\n", + "comm = MPI.COMM_WORLD\n", + "size = comm.Get_size()\n", + "rank = comm.Get_rank()\n", + "if rank == 0:\n", + " data = np.arange(10, dtype='i')\n", + " comm.Send([data, MPI.INT], dest=1, tag=77)\n", + "elif rank == 1:\n", + " data = np.empty(10, dtype='i')\n", + " comm.Recv([data, MPI.INT], source=0, tag=77)\n", + " print(data)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6e0c463e-72a6-4bd8-8c12-f80ed5831823", + "metadata": { + "id": "6e0c463e-72a6-4bd8-8c12-f80ed5831823" + }, + "outputs": [], + "source": [ + "!$SRUN -n 2 -c 2 python send_recv_eff.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "896b3bcd-d66b-49aa-9f16-4318c9fdb9fb", + "metadata": { + "id": "896b3bcd-d66b-49aa-9f16-4318c9fdb9fb" + }, + "outputs": [], + "source": [ + "%%writefile send_recv_numpy.py\n", + "from mpi4py import MPI\n", + "import numpy as np\n", + "\n", + "comm = MPI.COMM_WORLD\n", + "size = comm.Get_size()\n", + "rank = comm.Get_rank()\n", + "\n", + "if rank == 0:\n", + " data = np.arange(10, dtype=np.float64)\n", + " comm.Send(data, dest=1, tag=13)\n", + "elif rank == 1:\n", + " data = np.empty(10, dtype=np.float64)\n", + " comm.Recv(data, source=0, tag=13)\n", + " print(data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4ba204e8-ee72-4ef2-a04f-96de7a779e2b", + "metadata": { + "id": "4ba204e8-ee72-4ef2-a04f-96de7a779e2b" + }, + "outputs": [], + "source": [ + "!$SRUN -n 2 -c 2 python send_recv_numpy.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4107702e-c9bf-4917-a5e0-4b87a505fc66", + "metadata": { + "id": "4107702e-c9bf-4917-a5e0-4b87a505fc66" + }, + "outputs": [], + "source": [ + "%%writefile barrier.py\n", + "\n", + "from mpi4py import MPI\n", + "comm = MPI.COMM_WORLD\n", + "rank = comm.Get_rank()\n", + "for r_id in range(comm.Get_size()):\n", + " if rank == r_id:\n", + " print(f\"Hello from proc:\")\n", + "comm.Barrier()\n", + "print(f\"proc {rank} after barrier\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a10c0baf-b6a6-415e-8083-ba7ade5273a4", + "metadata": { + "id": "a10c0baf-b6a6-415e-8083-ba7ade5273a4" + }, + "outputs": [], + "source": [ + "!$SRUN -n 4 -c 4 python barrier.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "711866b0-b432-4252-80c2-25d722bab297", + "metadata": { + "id": "711866b0-b432-4252-80c2-25d722bab297" + }, + "outputs": [], + "source": [ + "%%writefile scatter.py\n", + "from mpi4py import MPI\n", + "comm = MPI.COMM_WORLD\n", + "rank = comm.Get_rank()\n", + "size = comm.Get_size()\n", + "\n", + "if rank == 0:\n", + " data = [(i+1)**2 for i in range(size)]\n", + "else:\n", + " data = None\n", + "\n", + "data = comm.scatter(data, root=0)\n", + "assert data == (rank+1)**2\n", + "\n", + "print(f\"data on rank %d is: {rank}, {data}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d72dd770-b215-4eed-8f2a-86e25f6c2c93", + "metadata": { + "id": "d72dd770-b215-4eed-8f2a-86e25f6c2c93" + }, + "outputs": [], + "source": [ + "!$SRUN -n 4 -c 4 python scatter.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "37422960-07af-4bde-9191-3a7e5c086740", + "metadata": { + "id": "37422960-07af-4bde-9191-3a7e5c086740" + }, + "outputs": [], + "source": [ + "%%writefile gather.py\n", + "from mpi4py import MPI\n", + "comm = MPI.COMM_WORLD\n", + "rank = comm.Get_rank()\n", + "size = comm.Get_size()\n", + "\n", + "data = (rank+1)**2\n", + "print(f\"before gather, data on rank %d is: {rank}, {data}\")\n", + "\n", + "comm.Barrier()\n", + "data = comm.gather(data, root=0)\n", + "if rank == 0:\n", + " for i in range(size):\n", + " assert data[i] == (i+1)**2\n", + "else:\n", + " assert data is None\n", + "\n", + "print(f\"data on rank: {rank} is: {data}\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "45ef2e24-2634-4243-974e-2fc5e2719a15", + "metadata": { + "id": "45ef2e24-2634-4243-974e-2fc5e2719a15" + }, + "outputs": [], + "source": [ + "!$SRUN -n 4 -c 4 python gather.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b1b25c4b-1331-4c1e-9d71-67cb820af2f4", + "metadata": { + "id": "b1b25c4b-1331-4c1e-9d71-67cb820af2f4" + }, + "outputs": [], + "source": [ + "%%writefile reduce.py\n", + "from mpi4py import MPI\n", + "comm = MPI.COMM_WORLD\n", + "rank = comm.Get_rank()\n", + "size = comm.Get_size()\n", + "\n", + "sendmsg = rank\n", + "recvmsg1 = comm.reduce(sendmsg, op=MPI.SUM, root=0)\n", + "recvmsg2 = comm.allreduce(sendmsg)\n", + "\n", + "print(recvmsg2)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "44f164e9-d833-4a7a-8972-06a63ccae1a3", + "metadata": { + "id": "44f164e9-d833-4a7a-8972-06a63ccae1a3" + }, + "outputs": [], + "source": [ + "!$SRUN -n 4 -c 4 python reduce.py" + ] + }, + { + "cell_type": "markdown", + "id": "20a0b643-481c-4d61-a8a6-0320d977a69c", + "metadata": { + "id": "20a0b643-481c-4d61-a8a6-0320d977a69c" + }, + "source": [ + "## Compute Pi Example\n", + "\n", + "The following example is completely self-contained to simplify reuse in another script. You can switch between running the code in parallel and serial by executing an %autopx cell." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "86b7cdb2-f430-4b5b-a24f-ced21cad2a37", + "metadata": { + "id": "86b7cdb2-f430-4b5b-a24f-ced21cad2a37" + }, + "outputs": [], + "source": [ + "%%writefile compute_pi.py\n", + "from mpi4py import MPI\n", + "import math\n", + "\n", + "def compute_pi(n, start=0, step=1):\n", + " h = 1.0 / n\n", + " s = 0.0\n", + " for i in range(start, n, step):\n", + " x = h * (i + 0.5)\n", + " s += 4.0 / (1.0 + x**2)\n", + " return s * h\n", + "\n", + "comm = MPI.COMM_WORLD\n", + "nprocs = comm.Get_size()\n", + "myrank = comm.Get_rank()\n", + "if myrank == 0:\n", + " n = 10\n", + "else:\n", + " n = None\n", + "\n", + "n = comm.bcast(n, root=0)\n", + "\n", + "mypi = compute_pi(n, myrank, nprocs)\n", + "\n", + "pi = comm.reduce(mypi, op=MPI.SUM, root=0)\n", + "\n", + "if myrank == 0:\n", + " error = abs(pi - math.pi)\n", + " print(\"pi is approximately %.16f\\nerror is %.16f\" % (pi, error))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d4294276-7306-4a17-8315-85704c439a7c", + "metadata": { + "id": "d4294276-7306-4a17-8315-85704c439a7c" + }, + "outputs": [], + "source": [ + "!$SRUN -n 4 -c 4 python compute_pi.py" + ] + }, + { + "cell_type": "markdown", + "id": "57a66360-777b-4492-9094-324176a32659", + "metadata": { + "id": "57a66360-777b-4492-9094-324176a32659" + }, + "source": [ + "## Mandelbrot Set Example\n", + "\n", + "The following example is completely self-contained to simplify reuse in another script." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19ad96dd-f8fc-4a89-8d11-c50411944662", + "metadata": { + "id": "19ad96dd-f8fc-4a89-8d11-c50411944662" + }, + "outputs": [], + "source": [ + "%%writefile compute_mandelbot.py\n", + "import matplotlib.pyplot as plt\n", + "import matplotlib.cm as cm\n", + "\n", + "from mpi4py import MPI\n", + "import numpy as np\n", + "\n", + "def mandelbrot (x, y, maxit):\n", + " c = x + y*1j\n", + " z = 0 + 0j\n", + " it = 0\n", + " while abs(z) < 2 and it < maxit:\n", + " z = z**2 + c\n", + " it += 1\n", + " return it\n", + "\n", + "x1, x2 = -2.0, 1.0\n", + "y1, y2 = -1.0, 1.0\n", + "w, h = 250, 200\n", + "maxit = 127\n", + "\n", + "comm = MPI.COMM_WORLD\n", + "size = comm.Get_size()\n", + "rank = comm.Get_rank()\n", + "\n", + "# number of rows to compute here\n", + "N = h // size + (h % size > rank)\n", + "\n", + "# first row to compute here\n", + "start = comm.scan(N)-N\n", + "\n", + "# array to store local result\n", + "Cl = np.zeros([N, w], dtype='i')\n", + "\n", + "# compute owned rows\n", + "\n", + "dx = (x2 - x1) / w\n", + "dy = (y2 - y1) / h\n", + "\n", + "for i in range(N):\n", + " y = y1 + (i + start) * dy\n", + " for j in range(w):\n", + " x = x1 + j * dx\n", + " Cl[i, j] = mandelbrot(x, y, maxit)\n", + "\n", + "# gather results at root (process 0)\n", + "counts = comm.gather(N, root=0)\n", + "C = None\n", + "if rank == 0:\n", + " C = np.zeros([h, w], dtype='i')\n", + "\n", + "# here we create a custom datatype for sending/receiving rows of data.\n", + "rowtype = MPI.INT.Create_contiguous(w)\n", + "rowtype.Commit()\n", + "\n", + "comm.Gatherv(sendbuf=[Cl, MPI.INT], recvbuf=[C, (counts, None), rowtype],root=0)\n", + "\n", + "rowtype.Free()\n", + "\n", + "if rank==0:\n", + " # Do the plotting on rank 0\n", + " plt.figure()\n", + " plt.imshow(C, cmap=cm.Spectral)\n", + " plt.savefig('test.png')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ac4a537d-942a-4c09-a2db-21f6d0fc9d33", + "metadata": { + "id": "ac4a537d-942a-4c09-a2db-21f6d0fc9d33" + }, + "outputs": [], + "source": [ + "!$SRUN -n 4 -c 4 python compute_mandelbot.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1e83bd55-8a6d-4bbb-bd88-cdfd916690ab", + "metadata": { + "id": "1e83bd55-8a6d-4bbb-bd88-cdfd916690ab" + }, + "outputs": [], + "source": [ + "from IPython.display import Image\n", + "Image(filename='test.png')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6db131ed-12dd-4cfb-9101-5f8cf00a8eb3", + "metadata": { + "id": "6db131ed-12dd-4cfb-9101-5f8cf00a8eb3" + }, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "training-pyhpc-2025", + "language": "python", + "name": "pyhpc-2025" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.5" + }, + "colab": { + "provenance": [] + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file