{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quickstart: running JAX on IPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-05T17:39:19.669191Z",
     "iopub.status.busy": "2023-07-05T17:39:19.668950Z",
     "iopub.status.idle": "2023-07-05T17:39:39.776473Z",
     "shell.execute_reply": "2023-07-05T17:39:39.775703Z",
     "shell.execute_reply.started": "2023-07-05T17:39:19.669169Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found existing installation: jax 0.3.16+ipu\n",
      "Uninstalling jax-0.3.16+ipu:\n",
      "  Successfully uninstalled jax-0.3.16+ipu\n",
      "Found existing installation: jaxlib 0.3.15+ipu.sdk310\n",
      "Uninstalling jaxlib-0.3.15+ipu.sdk310:\n",
      "  Successfully uninstalled jaxlib-0.3.15+ipu.sdk310\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mLooking in links: https://graphcore-research.github.io/jax-experimental/wheels.html\n",
      "Collecting jax==0.3.16+ipu\n",
      "  Downloading https://github.com/graphcore-research/jax-experimental/releases/download/jax-v0.3.16-ipu-beta2-sdk3/jax-0.3.16%2Bipu-py3-none-any.whl (1.2 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.2/1.2 MB\u001b[0m \u001b[31m31.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting jaxlib==0.3.15+ipu.sdk310\n",
      "  Downloading https://github.com/graphcore-research/jax-experimental/releases/download/jax-v0.3.16-ipu-beta2-sdk3/jaxlib-0.3.15%2Bipu.sdk310-cp38-none-manylinux2014_x86_64.whl (109.4 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m109.4/109.4 MB\u001b[0m \u001b[31m30.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: etils[epath] in /usr/local/lib/python3.8/dist-packages (from jax==0.3.16+ipu) (1.3.0)\n",
      "Requirement already satisfied: absl-py in /usr/local/lib/python3.8/dist-packages (from jax==0.3.16+ipu) (1.4.0)\n",
      "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from jax==0.3.16+ipu) (4.5.0)\n",
      "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.8/dist-packages (from jax==0.3.16+ipu) (3.3.0)\n",
      "Requirement already satisfied: numpy>=1.20 in /usr/local/lib/python3.8/dist-packages (from jax==0.3.16+ipu) (1.24.4)\n",
      "Requirement already satisfied: pybind11 in /usr/local/lib/python3.8/dist-packages (from jax==0.3.16+ipu) (2.10.4)\n",
      "Requirement already satisfied: scipy>=1.5 in /usr/local/lib/python3.8/dist-packages (from jax==0.3.16+ipu) (1.10.1)\n",
      "Requirement already satisfied: cppimport in /usr/local/lib/python3.8/dist-packages (from jax==0.3.16+ipu) (22.8.2)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from cppimport->jax==0.3.16+ipu) (3.12.2)\n",
      "Requirement already satisfied: mako in /usr/local/lib/python3.8/dist-packages (from cppimport->jax==0.3.16+ipu) (1.2.4)\n",
      "Requirement already satisfied: zipp in /usr/local/lib/python3.8/dist-packages (from etils[epath]->jax==0.3.16+ipu) (3.14.0)\n",
      "Requirement already satisfied: importlib_resources in /usr/local/lib/python3.8/dist-packages (from etils[epath]->jax==0.3.16+ipu) (5.12.0)\n",
      "Requirement already satisfied: MarkupSafe>=0.9.2 in /usr/local/lib/python3.8/dist-packages (from mako->cppimport->jax==0.3.16+ipu) (2.1.2)\n",
      "Installing collected packages: jaxlib, jax\n",
      "Successfully installed jax-0.3.16+ipu jaxlib-0.3.15+ipu.sdk310\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "# Install experimental JAX for IPUs (SDK 3.1) from Github releases.\n",
    "!pip uninstall -y jax jaxlib\n",
    "!pip install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk310 -f https://graphcore-research.github.io/jax-experimental/wheels.html"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the number of IPUs (or use IPUModel emulator)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-05T17:39:39.777889Z",
     "iopub.status.busy": "2023-07-05T17:39:39.777713Z",
     "iopub.status.idle": "2023-07-05T17:39:41.231066Z",
     "shell.execute_reply": "2023-07-05T17:39:41.230321Z",
     "shell.execute_reply.started": "2023-07-05T17:39:39.777868Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Platform=ipu\n",
      "Number of devices=2\n",
      "IpuDevice(id=0, num_tiles=1472, version=ipu2)\n",
      "IpuDevice(id=1, num_tiles=1472, version=ipu2)\n"
     ]
    }
   ],
   "source": [
    "# Enable IPU compilation cache\n",
    "import os\n",
    "os.environ['TF_POPLAR_FLAGS'] = \"\"\"\n",
    "  --executable_cache_path=/tmp/ipu-ef-cache\n",
    "  --show_progress_bar=true\n",
    "\"\"\"\n",
    "\n",
    "import jax\n",
    "from jax.config import config\n",
    "\n",
    "# Uncomment to use IPU model emulator.\n",
    "# config.FLAGS.jax_ipu_use_model = True\n",
    "# config.FLAGS.jax_ipu_model_num_tiles = 8\n",
    "\n",
    "# Select how many IPUs will be visible.\n",
    "config.update('jax_ipu_device_count', 2)\n",
    "\n",
    "\n",
    "print(f\"Platform={jax.default_backend()}\")\n",
    "print(f\"Number of devices={jax.device_count()}\")\n",
    "devices = jax.devices()\n",
    "print(\"\\n\".join([str(d) for d in devices]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JAX basics on IPU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's run a simple jit function on single IPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-05T17:39:41.233281Z",
     "iopub.status.busy": "2023-07-05T17:39:41.232933Z",
     "iopub.status.idle": "2023-07-05T17:39:45.318527Z",
     "shell.execute_reply": "2023-07-05T17:39:45.317539Z",
     "shell.execute_reply.started": "2023-07-05T17:39:41.233258Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Result:\n",
      "[[ 2.2647033   0.03436255  1.272394  ]\n",
      " [ 0.85815954 -0.30243903 -0.06740052]\n",
      " [-0.87136644  0.67539465 -2.7730992 ]]\n",
      "Result.platform = ipu\n",
      "Result.device = IpuDevice(id=0, num_tiles=1472, version=ipu2)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from jax import jit\n",
    "import jax.numpy as jnp\n",
    "\n",
    "@jit\n",
    "def func(x, w, b):\n",
    "    return jnp.matmul(w, x) + b\n",
    "\n",
    "x = np.random.normal(size=[2, 3])\n",
    "w = np.random.normal(size=[3, 2])\n",
    "b = np.random.normal(size=[3, 3])\n",
    "\n",
    "r = func(x, w, b)\n",
    "print(f\"Result:\\n{r}\")\n",
    "print(f\"Result.platform = {r.platform()}\")\n",
    "print(f\"Result.device = {r.device()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the `jax.device_put` API, we can put variables onto specific devices. Here is an example to run the jit function on `ipu:1`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-05T17:39:45.320413Z",
     "iopub.status.busy": "2023-07-05T17:39:45.320211Z",
     "iopub.status.idle": "2023-07-05T17:39:49.292317Z",
     "shell.execute_reply": "2023-07-05T17:39:49.291614Z",
     "shell.execute_reply.started": "2023-07-05T17:39:45.320393Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Result:\n",
      "[[ 2.2647033   0.03436255  1.272394  ]\n",
      " [ 0.85815954 -0.30243903 -0.06740052]\n",
      " [-0.87136644  0.67539465 -2.7730992 ]]\n",
      "Result.platform = ipu\n",
      "Result.device = IpuDevice(id=1, num_tiles=1472, version=ipu2)\n"
     ]
    }
   ],
   "source": [
    "x = jax.device_put(x, devices[1])\n",
    "w = jax.device_put(w, devices[1])\n",
    "b = jax.device_put(b, devices[1])\n",
    "\n",
    "r = func(x, w, b)\n",
    "\n",
    "print(f\"Result:\\n{r}\")\n",
    "print(f\"Result.platform = {r.platform()}\")\n",
    "print(f\"Result.device = {r.device()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`jit` also allows us to choose the backend the function will be running on.  For example, the below function will be running on `cpu`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-05T17:39:49.294476Z",
     "iopub.status.busy": "2023-07-05T17:39:49.294208Z",
     "iopub.status.idle": "2023-07-05T17:39:49.351168Z",
     "shell.execute_reply": "2023-07-05T17:39:49.350347Z",
     "shell.execute_reply.started": "2023-07-05T17:39:49.294454Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Result:\n",
      "[[ 2.2647033   0.03436255  1.272394  ]\n",
      " [ 0.85815954 -0.30243903 -0.06740052]\n",
      " [-0.8713663   0.6753946  -2.7730992 ]]\n",
      "Result.platform = cpu\n",
      "Result.device = TFRT_CPU_0\n"
     ]
    }
   ],
   "source": [
    "from functools import partial\n",
    "\n",
    "@partial(jit, backend='cpu')\n",
    "def func(x, w, b):\n",
    "    return jnp.matmul(w, x) + b\n",
    "\n",
    "r = func(x, w, b)\n",
    "\n",
    "print(f\"Result:\\n{r}\")\n",
    "print(f\"Result.platform = {r.platform()}\")\n",
    "print(f\"Result.device = {r.device()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JAX Pseudo Random Numbers generation\n",
    "\n",
    "Reproducible random numbers across platforms using JAX ThreeFry PRNG.  We will run this on both CPU and IPU.\n",
    "\n",
    "This is a relatively complex workload for IPU, so the first time through it will take a few seconds to compile.\n",
    "Let's switch on logging to see the compilation in action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-05T17:39:49.352146Z",
     "iopub.status.busy": "2023-07-05T17:39:49.351981Z",
     "iopub.status.idle": "2023-07-05T17:39:49.357227Z",
     "shell.execute_reply": "2023-07-05T17:39:49.356625Z",
     "shell.execute_reply.started": "2023-07-05T17:39:49.352126Z"
    }
   },
   "outputs": [],
   "source": [
    "config.update('jax_log_compiles', True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-05T17:39:49.358067Z",
     "iopub.status.busy": "2023-07-05T17:39:49.357901Z",
     "iopub.status.idle": "2023-07-05T17:39:54.314171Z",
     "shell.execute_reply": "2023-07-05T17:39:54.313350Z",
     "shell.execute_reply.started": "2023-07-05T17:39:49.358049Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Finished tracing + transforming random_fn for jit in 0.015218019485473633 sec\n",
      "WARNING:absl:Compiling random_fn (140215381452256 for args (ShapedArray(int32[], weak_type=True),).\n",
      "WARNING:absl:Finished XLA compilation of random_fn in 0.08308672904968262 sec\n",
      "WARNING:absl:Finished tracing + transforming random_fn for jit in 0.0146026611328125 sec\n",
      "WARNING:absl:Compiling random_fn (140214644209120 for args (ShapedArray(int32[], weak_type=True),).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU PRNG: (DeviceArray([255383827, 267815257], dtype=uint32), DeviceArray([0.7367313 , 0.92771065, 0.91349196, 0.3181516 ], dtype=float32))\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Finished XLA compilation of random_fn in 4.008972644805908 sec\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IPU PRNG: (DeviceArray([255383827, 267815257], dtype=uint32), DeviceArray([0.7367313 , 0.92771065, 0.91349196, 0.3181516 ], dtype=float32))\n"
     ]
    }
   ],
   "source": [
    "def random_fn(seed: int):\n",
    "    key = jax.random.PRNGKey(seed)\n",
    "    k1, k2 = jax.random.split(key)\n",
    "    return k2, jax.random.uniform(k1, (4,))\n",
    "\n",
    "random_fn_cpu = jax.jit(random_fn, backend=\"cpu\")\n",
    "print(\"CPU PRNG:\", random_fn_cpu(42))\n",
    "\n",
    "random_fn_ipu = jax.jit(random_fn, backend=\"ipu\")\n",
    "print(\"IPU PRNG:\", random_fn_ipu(42))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you run the same cell again, JAX compilation will be triggered by the redefinition of `random_fn`, but the IPU compilation will hit the cache, and be super quick."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JAX asynchronous dispatch on IPUs\n",
    "\n",
    "JAX IPU supports synchronous dispatch, allowing simple and efficient implementation of:\n",
    "* Inference and training pipeline (see MNIST examples);\n",
    "* Pipelining between multiple IPUs;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-05T17:39:54.316536Z",
     "iopub.status.busy": "2023-07-05T17:39:54.316356Z",
     "iopub.status.idle": "2023-07-05T17:39:54.319848Z",
     "shell.execute_reply": "2023-07-05T17:39:54.319276Z",
     "shell.execute_reply.started": "2023-07-05T17:39:54.316516Z"
    }
   },
   "outputs": [],
   "source": [
    "@partial(jit, backend='ipu')\n",
    "def compute_fn(x, w):\n",
    "    return jnp.matmul(w, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-05T17:39:54.320898Z",
     "iopub.status.busy": "2023-07-05T17:39:54.320739Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Finished tracing + transforming compute_fn for jit in 0.0010917186737060547 sec\n",
      "WARNING:absl:Compiling compute_fn (140215852776880 for args (ShapedArray(float32[1024,1024]), ShapedArray(float32[1024,1024])).\n",
      "WARNING:absl:Finished XLA compilation of compute_fn in 10.512851238250732 sec\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray([[ 14.617263 , -50.7595   ,  20.368622 , ...,  -6.4952183,\n",
       "               34.855034 , -54.406925 ],\n",
       "             [ 19.57885  , -12.667716 , -52.3782   , ..., -16.635777 ,\n",
       "               16.096308 , -27.73622  ],\n",
       "             [ 25.043592 , -12.324875 ,  24.267931 , ...,  19.963419 ,\n",
       "                1.868267 ,  32.6991   ],\n",
       "             ...,\n",
       "             [-32.77336  ,  -2.1762786,  33.599766 , ..., -69.713974 ,\n",
       "               16.669662 ,  17.313225 ],\n",
       "             [ 36.69899  , -12.384048 ,  16.702576 , ..., -40.15704  ,\n",
       "              -14.347113 ,  12.609639 ],\n",
       "             [-58.898438 ,   7.7527933,  25.596447 , ..., -64.197975 ,\n",
       "              -22.680655 ,  11.010284 ]], dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = np.random.normal(size=[1024, 1024]).astype(np.float32)\n",
    "# First run to compile jitted function, and load it on IPU.\n",
    "compute_fn(x, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-05T17:40:06.635211Z",
     "iopub.status.busy": "2023-07-05T17:40:06.635025Z",
     "iopub.status.idle": "2023-07-05T17:40:06.647642Z",
     "shell.execute_reply": "2023-07-05T17:40:06.647053Z",
     "shell.execute_reply.started": "2023-07-05T17:40:06.635190Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 470 µs, sys: 94 µs, total: 564 µs\n",
      "Wall time: 410 µs\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "DeviceArray([[ 14.617263 , -50.7595   ,  20.368622 , ...,  -6.4952183,\n",
       "               34.855034 , -54.406925 ],\n",
       "             [ 19.57885  , -12.667716 , -52.3782   , ..., -16.635777 ,\n",
       "               16.096308 , -27.73622  ],\n",
       "             [ 25.043592 , -12.324875 ,  24.267931 , ...,  19.963419 ,\n",
       "                1.868267 ,  32.6991   ],\n",
       "             ...,\n",
       "             [-32.77336  ,  -2.1762786,  33.599766 , ..., -69.713974 ,\n",
       "               16.669662 ,  17.313225 ],\n",
       "             [ 36.69899  , -12.384048 ,  16.702576 , ..., -40.15704  ,\n",
       "              -14.347113 ,  12.609639 ],\n",
       "             [-58.898438 ,   7.7527933,  25.596447 , ..., -64.197975 ,\n",
       "              -22.680655 ,  11.010284 ]], dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# No blocking: benchmarking only dispatch time.\n",
    "%time w = compute_fn(x, x)\n",
    "w.block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-05T17:40:06.648601Z",
     "iopub.status.busy": "2023-07-05T17:40:06.648423Z",
     "iopub.status.idle": "2023-07-05T17:40:06.659853Z",
     "shell.execute_reply": "2023-07-05T17:40:06.658901Z",
     "shell.execute_reply.started": "2023-07-05T17:40:06.648582Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7.97 ms, sys: 0 ns, total: 7.97 ms\n",
      "Wall time: 7.18 ms\n"
     ]
    }
   ],
   "source": [
    "# Blocking: benchmarking properly the matmul.\n",
    "%time w = compute_fn(x, x).block_until_ready()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}