Norms.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f24a0770",
   "metadata": {},
   "source": [
    "I'd prefer columns over 3d arrays.  Is the norm just as fast?  Yes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9402b35b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CuPy version 11.0.0\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import cupy as cp\n",
    "from cupyx import jit\n",
    "print(f\"CuPy version {cp.__version__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cc06c844",
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 100000000\n",
    "mempool = cp.get_default_memory_pool()\n",
    "mempool.free_all_blocks()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bea73883",
   "metadata": {},
   "source": [
    "Using linalg.norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2eaf0409",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mempool.used_bytes 1200000000\n"
     ]
    }
   ],
   "source": [
    "x_v = cp.random.random((size, 3), dtype=np.float32)\n",
    "print(f\"mempool.used_bytes {mempool.used_bytes()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e58d44fb",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 80.9 ms, sys: 8.08 ms, total: 88.9 ms\n",
      "Wall time: 89.7 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "norm = cp.linalg.norm(x_v, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "35b40ac5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mempool.used_bytes 1600000000\n",
      "mempool.used_bytes 1200000000\n"
     ]
    }
   ],
   "source": [
    "print(f\"mempool.used_bytes {mempool.used_bytes()}\")\n",
    "del norm\n",
    "mempool.free_all_blocks()\n",
    "print(f\"mempool.used_bytes {mempool.used_bytes()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c684a89f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.02 ms, sys: 334 µs, total: 2.35 ms\n",
      "Wall time: 24.1 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "norm = cp.linalg.norm(x_v, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f2c55184",
   "metadata": {},
   "outputs": [],
   "source": [
    "del norm, x_v"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16434cb3",
   "metadata": {},
   "source": [
    "using cp.sqrt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e1a4c31b",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_col = cp.random.random(size, dtype=np.float32)\n",
    "y_col = cp.random.random(size, dtype=np.float32)\n",
    "z_col = cp.random.random(size, dtype=np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8359e669",
   "metadata": {},
   "outputs": [],
   "source": [
    "@cp.fuse()\n",
    "def norm_from_cols(x,y,z):\n",
    "    return cp.sqrt(x * x + y * y + z * z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "91393011",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 633 ms, sys: 144 µs, total: 633 ms\n",
      "Wall time: 633 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "norm_col = norm_from_cols(x_col, y_col, z_col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7b27f69f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mempool.used_bytes 1600000000\n",
      "mempool.used_bytes 1200000000\n"
     ]
    }
   ],
   "source": [
    "mempool.free_all_blocks()\n",
    "print(f\"mempool.used_bytes {mempool.used_bytes()}\")\n",
    "del norm_col\n",
    "mempool.free_all_blocks()\n",
    "print(f\"mempool.used_bytes {mempool.used_bytes()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "dbe40ec3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 74 µs, sys: 0 ns, total: 74 µs\n",
      "Wall time: 76.3 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "norm_col = norm_from_cols(x_col, y_col, z_col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8839361a",
   "metadata": {},
   "outputs": [],
   "source": [
    "del norm_col"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70e3e4f8",
   "metadata": {},
   "source": [
    "using rawkernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6b647856",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/joel/.local/lib/python3.8/site-packages/cupyx/jit/_interface.py:161: FutureWarning: cupyx.jit.rawkernel is experimental. The interface can change in the future.\n",
      "  cupy._util.experimental('cupyx.jit.rawkernel')\n"
     ]
    }
   ],
   "source": [
    "@jit.rawkernel()\n",
    "def normalize_in_place(x, y, z, size):\n",
    "    tid = jit.blockIdx.x * jit.blockDim.x + jit.threadIdx.x\n",
    "    ntid = jit.gridDim.x * jit.blockDim.x\n",
    "    for i in range(tid, size, ntid):\n",
    "        norm = cp.sqrt(x[i] * x[i] + y[i] * y[i] + z[i] * z[i])\n",
    "        x[i] = x[i] / norm\n",
    "        y[i] = y[i] / norm\n",
    "        z[i] = z[i] / norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a8b0f874",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 88.5 ms, sys: 85 µs, total: 88.6 ms\n",
      "Wall time: 88.3 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "normalize_in_place((128,), (1024,), (x_col, y_col, z_col, size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3211053e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 83 µs, sys: 0 ns, total: 83 µs\n",
      "Wall time: 85.4 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "normalize_in_place((128,), (1024,), (x_col, y_col, z_col, size))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "edc3af27",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mempool.used_bytes 1200000000\n",
      "mempool.used_bytes 0\n"
     ]
    }
   ],
   "source": [
    "mempool.free_all_blocks()\n",
    "print(f\"mempool.used_bytes {mempool.used_bytes()}\")\n",
    "del x_col, y_col, z_col\n",
    "mempool.free_all_blocks()\n",
    "print(f\"mempool.used_bytes {mempool.used_bytes()}\")"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}