{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0cb78f89-aa0d-4f34-97d7-048c2b0d9667",
   "metadata": {},
   "source": [
    "# Host vs. Device vs. Global"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "977c8350-f2f3-420d-b148-b49925103894",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numba\n",
    "from numba import cuda\n",
    "import numpy as np\n",
    "\n",
    "#no decorator = host function. Normal python function which runs on the CPU\n",
    "def host_func():\n",
    "    #create a numpy array of integers, x_host\n",
    "    x_host = np.zeros(100, dtype=int)\n",
    "\n",
    "    #create a copy of x_host on the GPU. This copy is in device memory, and cannot be directly accessed by the CPU\n",
    "    x_device = cuda.to_device(x_host)\n",
    "\n",
    "    #here, we call the global function, from the CPU, to be executed on the GPU. We pass the device array, x_device, to this function running on the GPU\n",
    "    #note the bracket notation, where we tell the GPU how many threads, and how many blocks of threads, to use\n",
    "    thread_blocks = 256\n",
    "    threads_per_block = 256\n",
    "    global_func[thread_blocks, threads_per_block](x_device)\n",
    "\n",
    "    #finally, copy the results, which have ben stored in x_device, back to x_host\n",
    "    x_device.copy_to_host(x_host)\n",
    "\n",
    "    #print the results (on the CPU!)\n",
    "    print(x_host)\n",
    "\n",
    "#this decorator tells numba to rewrite this python function as a global function to be run on the GPU, whenever it's called by the CPU\n",
    "@cuda.jit()\n",
    "def global_func(x):\n",
    "    thread_id = cuda.grid(1)\n",
    "    num_threads = cuda.gridsize(1)\n",
    "\n",
    "    #this for loop ensures each thread only works on part of the array, thread 0 works on array index 0, thread 1 works on array index 1, etc.\n",
    "    #extra threads may go unused, which is completely ok other than being a little inefficient\n",
    "    #if the array was larger than num_threads, threads would run on more than one index, due to the stride being equal to num_threads\n",
    "    #e.g. if num_threads = 16, thread 0 would work on array indices 0, 16, 32, 48, 64, 80, and 96 for this array of size 100\n",
    "    for i in range(thread_id, x.shape[0], num_threads):\n",
    "        x[i] += device_func(i)\n",
    "\n",
    "#this decorator, slightly different than the previous one, defines this function as a device function,\n",
    "#    which may only be called on the GPU (by a global function, or another device function)\n",
    "@cuda.jit(device=True)\n",
    "def device_func(value):\n",
    "    return value*value "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9af78cd6-ae99-4348-8491-08f259cff4dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[   0    1    4    9   16   25   36   49   64   81  100  121  144  169\n",
      "  196  225  256  289  324  361  400  441  484  529  576  625  676  729\n",
      "  784  841  900  961 1024 1089 1156 1225 1296 1369 1444 1521 1600 1681\n",
      " 1764 1849 1936 2025 2116 2209 2304 2401 2500 2601 2704 2809 2916 3025\n",
      " 3136 3249 3364 3481 3600 3721 3844 3969 4096 4225 4356 4489 4624 4761\n",
      " 4900 5041 5184 5329 5476 5625 5776 5929 6084 6241 6400 6561 6724 6889\n",
      " 7056 7225 7396 7569 7744 7921 8100 8281 8464 8649 8836 9025 9216 9409\n",
      " 9604 9801]\n"
     ]
    }
   ],
   "source": [
    "host_func()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4faa04b8-f3b9-4627-89d3-226e9a3d2194",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy as cp\n",
    "from cupyx import jit\n",
    "import numpy as np\n",
    "\n",
    "#no decorator = host function. Normal python function which runs on the CPU\n",
    "def host_func():\n",
    "    #create a numpy array of integers, x_host\n",
    "    x_host = np.zeros(100, dtype=int)\n",
    "\n",
    "    #create a copy of x_host on the GPU. This copy is in device memory, and cannot be directly accessed by the CPU\n",
    "    x_device = cp.array(x_host)\n",
    "\n",
    "    #here, we call the global function, from the CPU, to be executed on the GPU. We pass the device array, x_device, to this function running on the GPU\n",
    "    #note the bracket notation, where we tell the GPU how many threads, and how many blocks of threads, to use\n",
    "    thread_blocks = 256\n",
    "    threads_per_block = 256\n",
    "    global_func[thread_blocks, threads_per_block](x_device)\n",
    "\n",
    "    #finally, copy the results, which have ben stored in x_device, back to x_host\n",
    "    x_host = x_device.get()\n",
    "\n",
    "    #print the results (on the CPU!)\n",
    "    print(x_host)\n",
    "\n",
    "#this decorator tells numba to rewrite this python function as a global function to be run on the GPU, whenever it's called by the CPU\n",
    "@jit.rawkernel()\n",
    "def global_func(x):\n",
    "    thread_id = jit.grid(1)\n",
    "    num_threads = jit.gridsize(1)\n",
    "\n",
    "    #this for loop ensures each thread only works on part of the array, thread 0 works on array index 0, thread 1 works on array index 1, etc.\n",
    "    #extra threads may go unused, which is completely ok other than being a little inefficient\n",
    "    #if the array was larger than num_threads, threads would run on more than one index, due to the stride being equal to num_threads\n",
    "    #e.g. if num_threads = 16, thread 0 would work on array indices 0, 16, 32, 48, 64, 80, and 96 for this array of size 100\n",
    "    for i in range(thread_id, x.shape[0], num_threads):\n",
    "        x[i] += device_func(i)\n",
    "\n",
    "#this decorator, slightly different than the previous one, defines this function as a device function,\n",
    "#    which may only be called on the GPU (by a global function, or another device function)\n",
    "@jit.rawkernel(device=True)\n",
    "def device_func(value):\n",
    "    return value*value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "479feb4e-ac0e-4314-88ca-41f9a4707aa7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[   0    1    4    9   16   25   36   49   64   81  100  121  144  169\n",
      "  196  225  256  289  324  361  400  441  484  529  576  625  676  729\n",
      "  784  841  900  961 1024 1089 1156 1225 1296 1369 1444 1521 1600 1681\n",
      " 1764 1849 1936 2025 2116 2209 2304 2401 2500 2601 2704 2809 2916 3025\n",
      " 3136 3249 3364 3481 3600 3721 3844 3969 4096 4225 4356 4489 4624 4761\n",
      " 4900 5041 5184 5329 5476 5625 5776 5929 6084 6241 6400 6561 6724 6889\n",
      " 7056 7225 7396 7569 7744 7921 8100 8281 8464 8649 8836 9025 9216 9409\n",
      " 9604 9801]\n"
     ]
    }
   ],
   "source": [
    "host_func()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87b5cc7d-3544-42ed-a679-7154f4ed3052",
   "metadata": {},
   "source": [
    "# Detailed breakdown of the for loop - parallelizing the problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2541230b-0830-44bc-9f18-205e00f24106",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Case 1: num_tasks >= num_threads\n",
      "range(0, 8, 4) = [0, 4]\n",
      "range(1, 8, 4) = [1, 5]\n",
      "range(2, 8, 4) = [2, 6]\n",
      "range(3, 8, 4) = [3, 7]\n",
      "\n",
      "Case 3: num_tasks < num_threads\n",
      "range(0, 8, 100) = [0]\n",
      "range(1, 8, 100) = [1]\n",
      "...\n",
      "range(7, 8, 100) = [7]\n",
      "range(8, 8, 100) = []\n",
      "range(9, 8, 100) = []\n",
      "...\n",
      "range(99, 8, 100) = []\n"
     ]
    }
   ],
   "source": [
    "# print each of the ranges mentioned in the section. Convert each range to a list to print the actual values!\n",
    "\n",
    "#ranges are in the format \"range(thread_id, num_tasks, num_threads)\"\n",
    "#two cases of interest, where num_tasks >= num_threads, and where num_tasks < num_threads\n",
    "\n",
    "#First, the ranges when num_tasks is at least as big as num_threads\n",
    "print(\"Case 1: num_tasks >= num_threads\")\n",
    "print(f\"range(0, 8, 4) = {list(range(0, 8, 4))}\")\n",
    "print(f\"range(1, 8, 4) = {list(range(1, 8, 4))}\")\n",
    "print(f\"range(2, 8, 4) = {list(range(2, 8, 4))}\")\n",
    "print(f\"range(3, 8, 4) = {list(range(3, 8, 4))}\")\n",
    "print(\"\")\n",
    "\n",
    "#Second, the ranges when num_tasks is smaller than num_threads\n",
    "print(\"Case 3: num_tasks < num_threads\")\n",
    "print(f\"range(0, 8, 100) = {list(range(0, 8, 100))}\")\n",
    "print(f\"range(1, 8, 100) = {list(range(1, 8, 100))}\")\n",
    "print(\"...\")\n",
    "print(f\"range(7, 8, 100) = {list(range(7, 8, 100))}\")\n",
    "print(f\"range(8, 8, 100) = {list(range(8, 8, 100))}\")\n",
    "print(f\"range(9, 8, 100) = {list(range(9, 8, 100))}\")\n",
    "print(\"...\")\n",
    "print(f\"range(99, 8, 100) = {list(range(99, 8, 100))}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c7fa7669-21e4-4cd1-8596-34e9a36ec295",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2 2 2 2 2 2 2 2]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/g/g20/peters42/lassen_conda/envs/lassen/lib/python3.11/site-packages/numba/cuda/dispatcher.py:536: NumbaPerformanceWarning: Grid size 1 will likely result in GPU under-utilization due to low occupancy.\n",
      "  warn(NumbaPerformanceWarning(msg))\n"
     ]
    }
   ],
   "source": [
    "import numba\n",
    "from numba import cuda\n",
    "import numpy as np\n",
    "\n",
    "#all these kernels define an incrementation kernel, where the value at each array location is increased by 1\n",
    "#note that case1 and case3 are identical, the only difference comes from defining the number of threads!\n",
    "\n",
    "arr_host = np.ones(8, dtype=int)\n",
    "arr_device = cuda.to_device(arr_host)\n",
    "\n",
    "@cuda.jit()\n",
    "def case1(arr):\n",
    "    thread_id = cuda.grid(1)\n",
    "    num_threads = cuda.gridsize(1)\n",
    "    num_tasks = arr.shape[0]\n",
    "\n",
    "    for i in range(thread_id, num_tasks, num_threads):\n",
    "        arr[i] += 1\n",
    "\n",
    "case1[1, 4](arr_device) # 1x4 = 4 total threads\n",
    "\n",
    "arr_host = arr_device.copy_to_host()\n",
    "print(arr_host)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "99d7a37b-39b9-48f4-a580-603b55969e32",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[2 2 2 2]\n",
      " [2 2 2 2]\n",
      " [2 2 2 2]\n",
      " [2 2 2 2]]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/g/g20/peters42/lassen_conda/envs/lassen/lib/python3.11/site-packages/numba/cuda/dispatcher.py:536: NumbaPerformanceWarning: Grid size 1 will likely result in GPU under-utilization due to low occupancy.\n",
      "  warn(NumbaPerformanceWarning(msg))\n"
     ]
    }
   ],
   "source": [
    "import numba\n",
    "from numba import cuda\n",
    "import numpy as np\n",
    "\n",
    "#all these kernels define an incrementation kernel, where the value at each array location is increased by 1\n",
    "\n",
    "arr_host = np.ones([4, 4], dtype=int)\n",
    "arr_device = cuda.to_device(arr_host)\n",
    "\n",
    "@cuda.jit()\n",
    "def case2(arr):\n",
    "    thread_id_y, thread_id_x = cuda.grid(2)\n",
    "    num_threads_y, num_threads_x = cuda.gridsize(2)\n",
    "    num_tasks_y = arr.shape[0]\n",
    "    num_tasks_x = arr.shape[1]\n",
    "    # above two lines can alternatively be written as the line below:\n",
    "    #num_tasks_y, num_tasks_x = arr.shape\n",
    "\n",
    "    for i in range(thread_id_y, num_tasks_y, num_threads_y):\n",
    "        for j in range(thread_id_x, num_tasks_x, num_threads_x):\n",
    "            arr[i][j] += 1\n",
    "\n",
    "case2[1, 4](arr_device) # 1x4 = 4 total threads, implicitly reshaped into a 2x2 array of threads in the kernel\n",
    "\n",
    "arr_host = arr_device.copy_to_host()\n",
    "print(arr_host)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a4e1df20-f6ac-49df-97a5-60f12e77b953",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2 2 2 2 2 2 2 2]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/g/g20/peters42/lassen_conda/envs/lassen/lib/python3.11/site-packages/numba/cuda/dispatcher.py:536: NumbaPerformanceWarning: Grid size 1 will likely result in GPU under-utilization due to low occupancy.\n",
      "  warn(NumbaPerformanceWarning(msg))\n"
     ]
    }
   ],
   "source": [
    "import numba\n",
    "from numba import cuda\n",
    "import numpy as np\n",
    "\n",
    "#all these kernels define an incrementation kernel, where the value at each array location is increased by 1\n",
    "#note that case1 and case3 are identical, the only difference comes from defining the number of threads!\n",
    "\n",
    "arr_host = np.ones(8, dtype=int)\n",
    "arr_device = cuda.to_device(arr_host)\n",
    "\n",
    "@cuda.jit()\n",
    "def case3(arr):\n",
    "    thread_id = cuda.grid(1)\n",
    "    num_threads = cuda.gridsize(1)\n",
    "    num_tasks = arr.shape[0]\n",
    "\n",
    "    for i in range(thread_id, num_tasks, num_threads):\n",
    "        arr[i] += 1\n",
    "\n",
    "case3[1, 100](arr_device) # 1x100 = 100 total threads\n",
    "\n",
    "arr_host = arr_device.copy_to_host()\n",
    "print(arr_host)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5e53a8ea-2dda-4a87-8fd6-b10b56a18e53",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 16472 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 38992 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from MyDiffusionGPU import MyDiffusionClass\n",
    "\n",
    "sim = MyDiffusionClass(dimensions=[500, 500])\n",
    "\n",
    "#initialize non-array parameters\n",
    "sim.set_framework(\"GPU_SERIAL\") #\"CPU_SERIAL\" or \"GPU_SERIAL\"\n",
    "sim.set_dt(0.1)\n",
    "sim.set_save_path(\"data/diffusion_test\")\n",
    "sim.set_autosave_flag(True)\n",
    "sim.set_autosave_save_images_flag(True)\n",
    "sim.set_autosave_rate(2000)\n",
    "sim.set_boundary_conditions(\"PERIODIC\")\n",
    "\n",
    "data = {\n",
    "    \"D\":1.\n",
    "}\n",
    "sim.set_user_data(data)\n",
    "\n",
    "#initialize simulation arrays, all parameter changes should be BEFORE this point!\n",
    "sim.initialize_engine()\n",
    "\n",
    "#change array data here, for custom simulations\n",
    "\"\"\"\n",
    "sim.fields[0].data[:] = 1.\n",
    "length = sim.dimensions[0]\n",
    "width = sim.dimensions[1]\n",
    "sim.fields[0].data[length // 4:3 * length // 4, width // 4:3 * width // 4] = 0.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "#initial conditions\n",
    "sim.plot_simulation()\n",
    "\n",
    "#run simulation\n",
    "sim.simulate(2000)\n",
    "\n",
    "#final conditions\n",
    "sim.plot_simulation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34700abe-f505-4ba8-8a2d-4d6a8bc29454",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Lassen Kernel 2.0",
   "language": "python",
   "name": "lassen-kernel"
  },
  "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.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
