Smooth_Smith_Waterman_PyTorch_vs_JAX.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Smooth Smith Waterman PyTorch vs JAX",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyMqrzF9NI5LPl8BaIEyXMvx",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/shu65/blog-jax-notebook/blob/main/Smooth_Smith_Waterman_PyTorch_vs_JAX.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yN2QwDu0Df-b",
        "outputId": "5e1dc69f-458b-43bd-9426-c322c4c9890c"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Tue Nov 16 11:52:33 2021       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 495.44       Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla K80           Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   37C    P8    27W / 149W |      0MiB / 11441MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TwmPBBo6Dl2U",
        "outputId": "9b16f065-cb9c-4d4c-dcf2-ef24327b2cff"
      },
      "source": [
        "!pip3 install torch==1.10.0+cu111 torchvision==0.11.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in links: https://download.pytorch.org/whl/torch_stable.html\n",
            "Requirement already satisfied: torch==1.10.0+cu111 in /usr/local/lib/python3.7/dist-packages (1.10.0+cu111)\n",
            "Requirement already satisfied: torchvision==0.11.1+cu111 in /usr/local/lib/python3.7/dist-packages (0.11.1+cu111)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch==1.10.0+cu111) (3.10.0.2)\n",
            "Requirement already satisfied: pillow!=8.3.0,>=5.3.0 in /usr/local/lib/python3.7/dist-packages (from torchvision==0.11.1+cu111) (7.1.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torchvision==0.11.1+cu111) (1.19.5)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gcHHY-4MinUY",
        "outputId": "6e036a7e-83f9-4396-b803-33486d4eaac9"
      },
      "source": [
        "!pip3 list | grep -e jax -e torch"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "jax                           0.2.21\n",
            "jaxlib                        0.1.71+cuda111\n",
            "torch                         1.10.0+cu111\n",
            "torchsummary                  1.5.1\n",
            "torchtext                     0.11.0\n",
            "torchvision                   0.11.1+cu111\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xj_3Mo3shwfl"
      },
      "source": [
        "import time\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "import jax\n",
        "import jax.numpy as jnp"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SidHo-h5h-Ec"
      },
      "source": [
        "#torch_device = \"cpu\"\n",
        "torch_device = \"cuda\"\n",
        "\n",
        "np.random.seed(0)\n",
        "seq_1_len = 100\n",
        "seq_2_len = 150\n",
        "n_trials = 10\n",
        "\n",
        "score_matrix_np = np.random.random((seq_1_len, seq_2_len)).astype(\"float32\")\n",
        "score_matrix_torch = torch.as_tensor(score_matrix_np, device=torch_device)\n",
        "score_matrix_jnp = jax.device_put(jnp.array(score_matrix_np))                                                                                 \n",
        "\n",
        "seq_1_max_len = 100\n",
        "seq_2_max_len = 120\n",
        "num_pairs = 64\n",
        "\n",
        "batch_score_matrix_np = np.random.random((num_pairs, seq_1_len, seq_2_len)).astype(\"float32\")\n",
        "batch_lens_np = np.array([[np.random.choice([80,90,100]),np.random.choice([95,105,120])] for _ in range(num_pairs)])\n",
        "\n",
        "batch_score_matrix_torch = torch.as_tensor(batch_score_matrix_np, device=torch_device)\n",
        "batch_lens_torch = torch.as_tensor(batch_lens_np, device=torch_device)\n",
        "\n",
        "batch_score_matrix_jnp = jax.device_put(jnp.array(batch_score_matrix_np))                                                                     \n",
        "batch_lens_jnp = jax.device_put(jnp.array(batch_lens_np))    "
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zkL6O2MgiNrC",
        "outputId": "fbe53776-4f0d-45d9-92d5-6935a96d1436"
      },
      "source": [
        "def sw_np(NINF=-1e30):\n",
        "    \n",
        "    def _logsumexp(y, axis):\n",
        "        y = np.maximum(y,NINF)\n",
        "        return y.max(axis) + np.log(np.sum(np.exp(y - y.max(axis, keepdims=True)), axis=axis))\n",
        "\n",
        "    def _soft_maximum(x, temp, axis=None):\n",
        "        return temp*_logsumexp(x/temp, axis)\n",
        "\n",
        "    def _sw(score_matrix, lengths, gap=0, temp=1.0):\n",
        "        real_a, real_b = lengths\n",
        "        hij = np.full((real_a + 1, real_b + 1), fill_value=NINF, dtype=np.float32)\n",
        "        for i in range(real_a):\n",
        "            for j in range(real_b):\n",
        "                s = score_matrix[i, j]\n",
        "                m = hij[i, j] + s\n",
        "                g0 = hij[i + 1, j] + gap\n",
        "                g1 = hij[i, j + 1] + gap\n",
        "\n",
        "                h = np.stack([m, g0, g1, s], -1)\n",
        "                hij[i + 1, j + 1] = _soft_maximum(h, temp=temp, axis=-1)\n",
        "        hij = hij[1:, 1:]\n",
        "        score = _soft_maximum(hij, temp=temp)\n",
        "        return score\n",
        "    return _sw\n",
        "\n",
        "def batch_sw_np(NINF=-1e30):\n",
        "    def _batch_sw(batch_score_matrix, batch_lengths, gap=0, temp=1.0):\n",
        "        n_batches = batch_score_matrix.shape[0]\n",
        "        sw_func = sw_np(NINF=NINF)\n",
        "        ret = [sw_func(batch_score_matrix[i], batch_lengths[i], gap=gap, temp=temp) \n",
        "               for i in range(n_batches)]\n",
        "        return np.array(ret)\n",
        "    return _batch_sw\n",
        "\n",
        "my_sw_func = sw_np()\n",
        "start = time.time()\n",
        "for i in range(n_trials):\n",
        "  score = my_sw_func(score_matrix_np, (seq_1_len, seq_2_len))\n",
        "elapsed_time = time.time() - start\n",
        "print(score)\n",
        "print(\"avg numpy version:\", elapsed_time/n_trials, 'sec.')\n",
        "\n",
        "my_sw_func = batch_sw_np()\n",
        "start = time.time()\n",
        "for i in range(n_trials):\n",
        "  score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0)\n",
        "elapsed_time = time.time() - start\n",
        "print(score)\n",
        "print(\"avg numpy batch version:\", elapsed_time/n_trials, 'sec.')"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "232.3118133544922\n",
            "avg numpy version: 0.8752012252807617 sec.\n",
            "[ 96.91699982 101.68367767 100.91819    106.12583923  97.63957977\n",
            " 117.04177856 111.9108963   98.19909668 117.86167908 106.50817871\n",
            " 117.40112305 106.88807678 101.1057663   93.29880524 106.5920639\n",
            "  98.1852417   98.35207367 117.18545532 105.11823273  92.83053589\n",
            "  98.24658203 107.01963806 111.48445892  93.74310303 100.79042816\n",
            " 116.46691132 117.11985779 107.21204376 107.86577606 101.42388153\n",
            "  97.99030304 117.33204651 106.12135315 101.53005219  96.26170349\n",
            " 111.8204422  101.49585724 111.94686127 105.30028534  98.19165802\n",
            " 117.63220978 102.38180542  96.12507629 102.4134903  108.64068604\n",
            " 105.61355591  97.12931061  94.71160126  98.14660645  96.24068451\n",
            " 108.331604    97.41623688 101.38352966 111.96533966 111.80926514\n",
            " 104.25274658  96.34142303 104.95452118  99.25166321  98.32733154\n",
            " 109.10884857 108.92740631 103.15089417 111.78417206]\n",
            "avg numpy batch version: 34.46293420791626 sec.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bEMqjNY6iRIT",
        "outputId": "55f4da5e-1b27-4b53-bae7-53a77c96d441"
      },
      "source": [
        "def sw_jax(unroll=2, NINF=-1e30):\n",
        "    \n",
        "    def _make_mask(score_matrix, lengths):\n",
        "        a,b = score_matrix.shape\n",
        "        real_a, real_b = lengths\n",
        "        mask = (jnp.arange(a) < real_a)[:,None] * (jnp.arange(b) < real_b)[None,:]\n",
        "        return mask\n",
        "\n",
        "    def _rotate(score_matrix):\n",
        "        a,b = score_matrix.shape\n",
        "        n,m = (a+b-1),(a+b)//2\n",
        "        ar,br = jnp.arange(a)[::-1,None], jnp.arange(b)[None,:]\n",
        "        i,j = (br-ar)+(a-1),(ar+br)//2\n",
        "        rotated_score_matrix = jnp.full([n,m],NINF).at[i,j].set(score_matrix)\n",
        "        reverse_idx = (i, j)\n",
        "        return rotated_score_matrix, reverse_idx\n",
        "\n",
        "    def _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp):\n",
        "        def scan_f(prev, scan_xs):\n",
        "            h2, h1 = prev\n",
        "            h1_T = _get_prev_gap_cell_score(\n",
        "                scan_xs[\"gap_cell_condition\"],\n",
        "                jnp.pad(h1[:-1], [1,0], constant_values=(NINF,NINF)),\n",
        "                jnp.pad(h1[1:], [0,1], constant_values=(NINF,NINF)),\n",
        "            )\n",
        "            a = h2 + scan_xs[\"rotated_score_matrix\"]\n",
        "            g0 = h1 + gap\n",
        "            g1 = h1_T + gap\n",
        "            s = scan_xs[\"rotated_score_matrix\"]\n",
        "\n",
        "            h0 = jnp.stack([a, g0, g1, s], -1)\n",
        "            h0 = _soft_maximum(h0, temp, -1)\n",
        "            return (h1,h0), h0\n",
        "        \n",
        "        a,b = score_matrix.shape\n",
        "        n,m = rotated_score_matrix.shape\n",
        "\n",
        "        scan_xs = {\n",
        "            \"rotated_score_matrix\": rotated_score_matrix,\n",
        "            \"gap_cell_condition\": (jnp.arange(n)+a%2)%2\n",
        "        }\n",
        "        scan_init = (jnp.full(m, NINF), jnp.full(m, NINF))\n",
        "        return scan_f, scan_xs, scan_init\n",
        "\n",
        "    def _rotate_in_reverse(rotated_dp_matrix, reverse_idx):\n",
        "        return rotated_dp_matrix[reverse_idx]\n",
        "\n",
        "    def _logsumexp(y, axis):\n",
        "        y = jnp.maximum(y,NINF)\n",
        "        return jax.nn.logsumexp(y, axis=axis)\n",
        "\n",
        "    def _logsumexp_with_mask(y, axis, mask):\n",
        "        y = jnp.maximum(y,NINF)\n",
        "        return y.max(axis) + jnp.log(jnp.sum(mask * jnp.exp(y - y.max(axis, keepdims=True)), axis=axis))\n",
        "\n",
        "    def _soft_maximum(x, temp, axis=None):\n",
        "        return temp*_logsumexp(x/temp, axis)\n",
        "\n",
        "    def _soft_maximum_with_mask(x, temp, mask, axis=None):\n",
        "        return temp*_logsumexp_with_mask(x/temp, axis, mask)\n",
        "\n",
        "    def _get_prev_gap_cell_score(cond, true, false): \n",
        "        return cond*true + (1-cond)*false\n",
        "    \n",
        "    def _sw(score_matrix, lengths, gap=0, temp=1.0):\n",
        "        mask = _make_mask(score_matrix, lengths)\n",
        "        masked_score_matrix = score_matrix + NINF * (1 - mask)\n",
        "        rotated_score_matrix, reverse_idx = _rotate(masked_score_matrix)\n",
        "        scan_f, scan_xs, scan_init = _prepare_scan_inputs(score_matrix, rotated_score_matrix, gap, temp)\n",
        "        rotated_hij = jax.lax.scan(scan_f, scan_init, scan_xs, unroll=unroll)[-1]\n",
        "        hij = _rotate_in_reverse(rotated_hij, reverse_idx)\n",
        "        score = _soft_maximum_with_mask(hij, temp, mask=mask, axis=None)\n",
        "        return score\n",
        "    return _sw\n",
        "\n",
        "def batch_sw_jax(unroll=2, NINF=-1e30):\n",
        "    sw_func = sw_jax(unroll=unroll, NINF=NINF)\n",
        "    batch_sw_func = jax.vmap(sw_func, (0, 0, None, None))\n",
        "    return batch_sw_func\n",
        "\n",
        "my_sw_func = sw_jax()\n",
        "print(\"jax default first call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "start = time.time()\n",
        "for i in range(n_trials):\n",
        "  score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "elapsed_time = time.time() - start\n",
        "print(score)\n",
        "print(\"avg jax version:\", elapsed_time/n_trials, 'sec.')\n",
        "\n",
        "my_sw_func = jax.jit(sw_jax())\n",
        "print(\"jax jit first call\")\n",
        "%time score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "start = time.time()\n",
        "for i in range(n_trials):\n",
        "  score = my_sw_func(score_matrix_jnp, (seq_1_len, seq_2_len)).block_until_ready()\n",
        "elapsed_time = time.time() - start\n",
        "print(score)\n",
        "print(\"avg jax jit version:\", elapsed_time/n_trials, 'sec.')\n",
        "\n",
        "my_sw_func = batch_sw_jax()\n",
        "print(\"batch jax batch default first call\")\n",
        "%time score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0).block_until_ready()\n",
        "start = time.time()\n",
        "for i in range(n_trials):\n",
        "  score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0).block_until_ready()\n",
        "elapsed_time = time.time() - start\n",
        "print(score)\n",
        "print(\"avg jax batch version:\", elapsed_time/n_trials, 'sec.')\n",
        "\n",
        "my_sw_func = jax.jit(batch_sw_jax())\n",
        "print(\"batch jax jit default first call\")\n",
        "%time score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0).block_until_ready()\n",
        "start = time.time()\n",
        "for i in range(n_trials):\n",
        "  score = my_sw_func(batch_score_matrix_np, batch_lens_np, -1.0, 1.0).block_until_ready()\n",
        "elapsed_time = time.time() - start\n",
        "print(score)\n",
        "print(\"avg jax jit batch version:\", elapsed_time/n_trials, 'sec.')"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "jax default first call\n",
            "CPU times: user 1.42 s, sys: 56.5 ms, total: 1.47 s\n",
            "Wall time: 3.65 s\n",
            "232.31183\n",
            "avg jax version: 0.5170674562454224 sec.\n",
            "jax jit first call\n",
            "CPU times: user 1.06 s, sys: 13.6 ms, total: 1.07 s\n",
            "Wall time: 950 ms\n",
            "232.31183\n",
            "avg jax jit version: 0.008703041076660156 sec.\n",
            "batch jax batch default first call\n",
            "CPU times: user 2.19 s, sys: 37.9 ms, total: 2.22 s\n",
            "Wall time: 3.63 s\n",
            "[ 96.917    101.68368  100.918175 106.12583   97.63958  117.04178\n",
            " 111.9109    98.1991   117.86167  106.50818  117.40112  106.88807\n",
            " 101.10577   93.298805 106.592064  98.185234  98.35207  117.18546\n",
            " 105.11823   92.830536  98.24659  107.01963  111.48446   93.7431\n",
            " 100.79042  116.46691  117.119865 107.21204  107.865776 101.42388\n",
            "  97.9903   117.332054 106.12135  101.53005   96.2617   111.82044\n",
            " 101.49586  111.94686  105.300285  98.19166  117.63221  102.381805\n",
            "  96.125084 102.41349  108.64067  105.613556  97.12931   94.71161\n",
            "  98.14661   96.240685 108.331604  97.41624  101.38353  111.96533\n",
            " 111.809265 104.25275   96.34143  104.95452   99.25166   98.32733\n",
            " 109.10885  108.927414 103.150894 111.78417 ]\n",
            "avg jax batch version: 0.884079909324646 sec.\n",
            "batch jax jit default first call\n",
            "CPU times: user 1.85 s, sys: 16.1 ms, total: 1.87 s\n",
            "Wall time: 1.58 s\n",
            "[ 96.917    101.68368  100.918175 106.12583   97.63958  117.04178\n",
            " 111.9109    98.1991   117.86167  106.50818  117.40112  106.88807\n",
            " 101.10577   93.298805 106.592064  98.185234  98.35207  117.18546\n",
            " 105.11823   92.830536  98.24659  107.01963  111.48446   93.7431\n",
            " 100.79042  116.46691  117.119865 107.21204  107.865776 101.42388\n",
            "  97.9903   117.332054 106.12135  101.53005   96.2617   111.82044\n",
            " 101.49586  111.94686  105.300285  98.19166  117.63221  102.381805\n",
            "  96.125084 102.41349  108.64067  105.613556  97.12931   94.71161\n",
            "  98.14661   96.240685 108.331604  97.41624  101.38353  111.96533\n",
            " 111.809265 104.25275   96.34143  104.95452   99.25166   98.32733\n",
            " 109.10885  108.927414 103.150894 111.78417 ]\n",
            "avg jax jit batch version: 0.014213824272155761 sec.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ms0P2d_x1box",
        "outputId": "d9937154-fda5-43e2-bca5-3ffd3724a051"
      },
      "source": [
        "import torch.nn as nn\n",
        "\n",
        "\n",
        "class SwTorch(nn.Module):\n",
        "    def __init__(self, unroll=2, NINF=-1e30, device=\"cpu\"):\n",
        "        super(SwTorch, self).__init__()\n",
        "        self.unroll = unroll\n",
        "        self.NINF = torch.tensor(NINF, device=device)\n",
        "        self.device = device\n",
        "\n",
        "    def _make_mask(self, score_matrix, lengths):\n",
        "        a,b = score_matrix.shape\n",
        "        real_a = lengths[0]\n",
        "        real_b = lengths[1]\n",
        "        mask = (torch.arange(a, device=self.device) < real_a)[:,None] & (torch.arange(b, device=self.device) < real_b)[None,:]\n",
        "        return mask\n",
        "\n",
        "    def _rotate(self, score_matrix):\n",
        "        a,b = score_matrix.shape\n",
        "        n,m = (a+b-1),(a+b)//2\n",
        "        ar = torch.flip(torch.arange(a, device=self.device), [0])[:, None]\n",
        "        br = torch.arange(b, device=self.device)[None,:]\n",
        "        i,j = (br-ar)+(a-1),(ar+br)//2\n",
        "        rotated_score_matrix = torch.full([n,m], self.NINF, dtype=score_matrix.dtype, device=self.device)\n",
        "        rotated_score_matrix[i, j] = score_matrix\n",
        "        reverse_idx = (i, j)\n",
        "        return rotated_score_matrix, reverse_idx\n",
        "\n",
        "    def _step(self, prev, gap_cell_condition, rotated_score_matrix, gap, temp):\n",
        "        h2,h1 = prev   # previous two rows of scoring (hij) mtx\n",
        "        h1_T = self._get_prev_gap_cell_score(\n",
        "            gap_cell_condition,\n",
        "            torch.nn.functional.pad(h1[:-1], [1,0], value=self.NINF),\n",
        "            torch.nn.functional.pad(h1[1:], [0,1], value=self.NINF),\n",
        "        )\n",
        "      \n",
        "        a = h2 + rotated_score_matrix\n",
        "        g0 = h1 + gap\n",
        "        g1 = h1_T + gap\n",
        "        s = rotated_score_matrix\n",
        "        h0 = torch.stack([a, g0, g1, s], -1)\n",
        "        h0 = self._soft_maximum(h0, temp, -1)\n",
        "        return (h1,h0), h0\n",
        "\n",
        "    def _rotate_in_reverse(self, rotated_dp_matrix, reverse_idx):\n",
        "        return rotated_dp_matrix[reverse_idx]\n",
        "\n",
        "    def _logsumexp(self, y, axis):\n",
        "        y = torch.maximum(y,self.NINF)\n",
        "        return torch.logsumexp(y, axis=axis)\n",
        "\n",
        "    def _logsumexp_with_mask(self, y, axis, mask):\n",
        "        y = torch.maximum(y,self.NINF)\n",
        "        if axis is None:\n",
        "          return torch.max(y) + torch.log(torch.sum(mask * torch.exp(y - torch.max(y))))\n",
        "        else:\n",
        "          return torch.max(y, axis)[0] + torch.log(torch.sum(mask * torch.exp(y - torch.max(y, axis, keepdims=True)[0]), axis=axis))\n",
        "\n",
        "    def _soft_maximum(self, x, temp, axis=None):\n",
        "        return temp*self._logsumexp(x/temp, axis)\n",
        "\n",
        "    def _soft_maximum_with_mask(self, x, temp, mask, axis=None):\n",
        "        return temp*self._logsumexp_with_mask(x/temp, axis, mask)\n",
        "\n",
        "    def _get_prev_gap_cell_score(self, cond, true, false): \n",
        "        return cond*true + (1-cond)*false\n",
        "\n",
        "    def forward(self, score_matrix, lengths, gap=0, temp=1.0):\n",
        "      mask = self._make_mask(score_matrix, lengths)\n",
        "      masked_score_matrix = score_matrix + self.NINF * (~mask)\n",
        "      rotated_score_matrix, reverse_idx = self._rotate(masked_score_matrix)\n",
        "\n",
        "      a,b = score_matrix.shape\n",
        "      n,m = rotated_score_matrix.shape\n",
        "      gap_cell_condition = (torch.arange(n, device=self.device)+a%2)%2\n",
        "      prev = (torch.full((m,), self.NINF, device=self.device), torch.full((m,), self.NINF, device=self.device))\n",
        "      rotated_hij = [None for _ in range(n)]\n",
        "      for i in range(n):\n",
        "          prev, h = self._step(prev, gap_cell_condition[i], rotated_score_matrix[i], gap, temp)\n",
        "          rotated_hij[i] = h\n",
        "      rotated_hij = torch.stack(rotated_hij)\n",
        "      hij = self._rotate_in_reverse(rotated_hij, reverse_idx)\n",
        "      score = self._soft_maximum_with_mask(hij, temp, mask=mask, axis=None)\n",
        "      return score\n",
        "\n",
        "\n",
        "class BatchSwTorch(nn.Module):\n",
        "    def __init__(self, unroll=2, NINF=-1e30, device=\"cpu\"):\n",
        "        super(BatchSwTorch, self).__init__()\n",
        "        self.device = device\n",
        "        self.sw = SwTorch(unroll=unroll, NINF=NINF, device=device)\n",
        "\n",
        "    def forward(self, batch_score_matrix, batch_lengths, gap=0, temp=1.0):\n",
        "        n_batches = batch_score_matrix.shape[0]\n",
        "        ret = torch.empty((n_batches,), dtype=batch_score_matrix.dtype, device=self.device)\n",
        "        for i in range(n_batches):\n",
        "          ret[i] = self.sw(batch_score_matrix[i], batch_lengths[i], gap=gap, temp=temp) \n",
        "        return ret\n",
        "\n",
        "\n",
        "lengths = torch.as_tensor([seq_1_len, seq_2_len])\n",
        "sw_module = SwTorch(device=torch_device)\n",
        "score = sw_module(score_matrix_torch, lengths)\n",
        "if torch_device == \"cuda\":  \n",
        "  torch.cuda.synchronize()\n",
        "start = time.time()\n",
        "for i in range(n_trials):\n",
        "   score = sw_module(score_matrix_torch, lengths)\n",
        "if torch_device == \"cuda\":  \n",
        "  torch.cuda.synchronize()\n",
        "elapsed_time = time.time() - start\n",
        "print(score)\n",
        "print(\"avg torch version:\", elapsed_time/n_trials, 'sec.')\n",
        "\n",
        "batch_sw_module = BatchSwTorch(device=torch_device)\n",
        "score = batch_sw_module(batch_score_matrix_torch, batch_lens_torch, torch.tensor(-1.0, device=torch_device), torch.tensor(1.0, device=torch_device))\n",
        "if torch_device == \"cuda\":  \n",
        "  torch.cuda.synchronize()\n",
        "start = time.time()\n",
        "for i in range(n_trials):\n",
        "  score = batch_sw_module(batch_score_matrix_torch, batch_lens_torch, torch.tensor(-1.0, device=torch_device), torch.tensor(1.0, device=torch_device))\n",
        "if torch_device == \"cuda\":  \n",
        "  torch.cuda.synchronize()\n",
        "elapsed_time = time.time() - start\n",
        "print(score)\n",
        "print(\"avg torch batch version:\", elapsed_time/n_trials, 'sec.')"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py:23: UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor').\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor(232.3118, device='cuda:0')\n",
            "avg torch version: 0.1302793025970459 sec.\n",
            "tensor([ 96.9170, 101.6837, 100.9182, 106.1258,  97.6396, 117.0418, 111.9109,\n",
            "         98.1991, 117.8617, 106.5082, 117.4011, 106.8881, 101.1058,  93.2988,\n",
            "        106.5921,  98.1852,  98.3521, 117.1855, 105.1182,  92.8305,  98.2466,\n",
            "        107.0196, 111.4845,  93.7431, 100.7904, 116.4669, 117.1199, 107.2120,\n",
            "        107.8658, 101.4239,  97.9903, 117.3321, 106.1214, 101.5301,  96.2617,\n",
            "        111.8204, 101.4959, 111.9469, 105.3003,  98.1917, 117.6322, 102.3818,\n",
            "         96.1251, 102.4135, 108.6407, 105.6136,  97.1293,  94.7116,  98.1466,\n",
            "         96.2407, 108.3316,  97.4162, 101.3835, 111.9653, 111.8093, 104.2527,\n",
            "         96.3414, 104.9545,  99.2517,  98.3273, 109.1088, 108.9274, 103.1509,\n",
            "        111.7842], device='cuda:0')\n",
            "avg torch batch version: 7.885982155799866 sec.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iCKym4tk61Ih",
        "outputId": "733c3fb3-f581-410c-9996-e38714a78e8f"
      },
      "source": [
        "from typing import Tuple\n",
        "\n",
        "\n",
        "def _make_batch_mask(batch_score_matrix, batch_lengths):\n",
        "    a, b, batch_size = batch_score_matrix.shape\n",
        "    real_a = batch_lengths[:, 0]\n",
        "    real_b = batch_lengths[:, 1]\n",
        "    mask_a = torch.arange(a, device=batch_score_matrix.device)[:, None].repeat(1, batch_size) < real_a[None, :]\n",
        "    mask_b = torch.arange(b, device=batch_score_matrix.device)[:, None].repeat(1, batch_size) < real_b[None, :]\n",
        "    mask = mask_a[:, None] & mask_b[None, :]\n",
        "    return mask\n",
        "\n",
        "\n",
        "def _logsumexp(y: torch.Tensor, axis: int, NINF: torch.Tensor) -> torch.Tensor:\n",
        "    y = torch.maximum(y, NINF)\n",
        "    return torch.logsumexp(y, dim=axis)\n",
        "\n",
        "\n",
        "def _logsumexp_with_mask(y: torch.Tensor, axis: int, mask: torch.Tensor, NINF: torch.Tensor) -> torch.Tensor:\n",
        "    y = torch.maximum(y, NINF)\n",
        "    return torch.max(y, axis)[0] + torch.log(torch.sum(mask * torch.exp(y - torch.max(y, dim=axis, keepdim=True)[0]), dim=axis))\n",
        "\n",
        "\n",
        "def _soft_maximum(x: torch.Tensor, temp: torch.Tensor, axis: int, NINF: torch.Tensor) -> torch.Tensor:\n",
        "    return temp*_logsumexp(x/temp, axis=axis, NINF=NINF)\n",
        "\n",
        "\n",
        "def _soft_maximum_with_mask(x: torch.Tensor, temp: torch.Tensor, axis: int, mask: torch.Tensor, NINF: torch.Tensor) -> torch.Tensor:\n",
        "    return temp*_logsumexp_with_mask(x/temp, axis=axis, mask=mask, NINF=NINF)\n",
        "\n",
        "\n",
        "def _rotate(batch_score_matrix: torch.Tensor, NINF: torch.Tensor, rotated_batch_score_matrix: torch.Tensor) -> Tuple[torch.Tensor, torch.Tenso\\\n",
        "r, torch.Tensor]:\n",
        "    a, b, batch_size = batch_score_matrix.shape\n",
        "    n,m = (a+b-1),(a+b)//2\n",
        "    ar = torch.flip(torch.arange(a, device=batch_score_matrix.device), [0])[:, None]\n",
        "    br = torch.arange(b, device=batch_score_matrix.device)[None,:]\n",
        "    i,j = (br-ar)+(a-1),(ar+br)//2      \n",
        "    rotated_batch_score_matrix[:, :, :] = NINF\n",
        "    rotated_batch_score_matrix[i, j, :] = batch_score_matrix                                                                                         \n",
        "    return rotated_batch_score_matrix, i, j\n",
        "\n",
        "\n",
        "def _rotate_in_reverse(rotated_dp_matrix, i, j):                                                                                                  \n",
        "    return rotated_dp_matrix[i, j]\n",
        "\n",
        "\n",
        "def _get_prev_gap_cell_score(cond, true, false):\n",
        "    return cond*true + (1-cond)*false\n",
        "\n",
        "@torch.jit.script\n",
        "def _step(h2, h1, gap_cell_condition, rotated_score_matrix, gap, temp, NINF, prev_gap_cell_true, prev_gap_cell_false):\n",
        "    prev_gap_cell_true[1:, :] = h1[:-1, :]\n",
        "    prev_gap_cell_false[:-1, :] = h1[1:, :]\n",
        "    h1_T = _get_prev_gap_cell_score(\n",
        "        gap_cell_condition,\n",
        "        prev_gap_cell_true,\n",
        "        prev_gap_cell_false,\n",
        "    )\n",
        "    a = h2 + rotated_score_matrix\n",
        "    g0 = h1 + gap\n",
        "    g1 = h1_T + gap\n",
        "    s = rotated_score_matrix\n",
        "    h0 = torch.stack([a, g0, g1, s], -1)\n",
        "    h0 = _soft_maximum(h0, temp, axis=-1, NINF=NINF)\n",
        "    return h1, h0, h0\n",
        "\n",
        "\n",
        "@torch.jit.script\n",
        "def _step_loop(init_h1, init_h0, gap_cell_condition, rotated_batch_score_matrix, gap, temp, NINF, prev_gap_cell_true, prev_gap_cell_false):\n",
        "    n, _, _ = rotated_batch_score_matrix.shape\n",
        "    rotated_hij = torch.empty((n, init_h1.shape[0], init_h1.shape[1]), dtype=init_h1.dtype, device=init_h1.device)\n",
        "    h1 = init_h1\n",
        "    h0 = init_h0\n",
        "    h1[:, :] = NINF\n",
        "    h0[:, :] = NINF\n",
        "    for i in range(n):\n",
        "        h1, h0, h = _step(h1, h0, gap_cell_condition=gap_cell_condition[i], rotated_score_matrix=rotated_batch_score_matrix[i], gap=gap, temp=temp, NINF=NINF, prev_gap_cell_true=prev_gap_cell_true, prev_gap_cell_false=prev_gap_cell_false,)\n",
        "        rotated_hij[i] = h                                                                                                 \n",
        "    return rotated_hij\n",
        "\n",
        "@torch.jit.script\n",
        "def batch_sw_func(batch_score_matrix, batch_lengths, gap, temp, NINF, rotated_batch_score_matrix, init_h1, init_h0, prev_gap_cell_true, prev_gap_cell_false):\n",
        "    transposed_batch_score_matrix = batch_score_matrix.permute(1, 2, 0)\n",
        "    mask = _make_batch_mask(transposed_batch_score_matrix, batch_lengths)\n",
        "    masked_batch_score_matrix = transposed_batch_score_matrix + NINF * (~mask)\n",
        "    rotated_batch_score_matrix, reverse_idx_i, reverse_idx_j = _rotate(masked_batch_score_matrix, NINF=NINF, rotated_batch_score_matrix=rotated_batch_score_matrix)\n",
        "    a, b, batch_size = transposed_batch_score_matrix.shape\n",
        "    n, m, _ = rotated_batch_score_matrix.shape\n",
        "    gap_cell_condition = (torch.arange(n, device=rotated_batch_score_matrix.device)+a%2)%2                                           \n",
        "    rotated_hij = _step_loop(init_h1, init_h0, gap_cell_condition, rotated_batch_score_matrix, gap, temp, prev_gap_cell_true=prev_gap_cell_true, prev_gap_cell_false=prev_gap_cell_false, NINF=NINF)                                                                                                   \n",
        "    hij = _rotate_in_reverse(rotated_hij, reverse_idx_i, reverse_idx_j)\n",
        "    score = _soft_maximum_with_mask(hij.reshape(a*b,batch_size), temp=temp, mask=mask.reshape(a*b, batch_size), axis=0, NINF=NINF)\n",
        "    return score\n",
        "\n",
        "\n",
        "batch_size, a, b = batch_score_matrix_torch.shape\n",
        "n,m = (a+b-1),(a+b)//2\n",
        "\n",
        "NINF=torch.tensor(-1e30, device=torch_device)\n",
        "rotated_batch_score_matrix = torch.full([n, m, batch_size], NINF, dtype=batch_score_matrix_torch.dtype, device=batch_score_matrix_torch.device\\\n",
        ")\n",
        "init_h1 = torch.full((m, batch_size), NINF, device=rotated_batch_score_matrix.device)\n",
        "init_h0 = torch.full((m, batch_size), NINF, device=rotated_batch_score_matrix.device)\n",
        "\n",
        "prev_gap_cell_true = torch.full((m, batch_size), NINF, device=rotated_batch_score_matrix.device)\n",
        "prev_gap_cell_false = torch.full((m, batch_size), NINF, device=rotated_batch_score_matrix.device)\n",
        "example_inputs = (\n",
        "    batch_score_matrix_torch,\n",
        "    batch_lens_torch,\n",
        "    torch.tensor(-1.0, device=torch_device),\n",
        "    torch.tensor(1.0, device=torch_device),\n",
        "    NINF,\n",
        "    rotated_batch_score_matrix,\n",
        "    init_h1,\n",
        "    init_h0,\n",
        "    prev_gap_cell_true,\n",
        "    prev_gap_cell_false,\n",
        ")                                                                                                  \n",
        "batch_sw = torch.jit.trace(batch_sw_func, example_inputs)\n",
        "print(batch_sw.graph)\n",
        "#batch_sw = batch_sw_func\n",
        "if torch_device == \"cuda\":\n",
        "  torch.cuda.synchronize()\n",
        "score = batch_sw(*example_inputs)\n",
        "if torch_device == \"cuda\":\n",
        "  torch.cuda.synchronize()\n",
        "if torch_device == \"cuda\":\n",
        "  torch.cuda.synchronize()\n",
        "start = time.time()\n",
        "for i in range(n_trials):\n",
        "  score = batch_sw(*example_inputs)\n",
        "if torch_device == \"cuda\":\n",
        "  torch.cuda.synchronize()\n",
        "elapsed_time = time.time() - start\n",
        "print(score)\n",
        "print(\"avg torch.jit.trace batch:\", elapsed_time/n_trials, 'sec.')\n"
      ],
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "graph(%0 : Float(64, 100, 150, strides=[15000, 150, 1], requires_grad=0, device=cuda:0),\n",
            "      %1 : Long(64, 2, strides=[2, 1], requires_grad=0, device=cuda:0),\n",
            "      %2 : Float(requires_grad=0, device=cuda:0),\n",
            "      %3 : Float(requires_grad=0, device=cuda:0),\n",
            "      %4 : Float(requires_grad=0, device=cuda:0),\n",
            "      %5 : Float(249, 125, 64, strides=[8000, 64, 1], requires_grad=0, device=cuda:0),\n",
            "      %6 : Float(125, 64, strides=[64, 1], requires_grad=0, device=cuda:0),\n",
            "      %7 : Float(125, 64, strides=[64, 1], requires_grad=0, device=cuda:0),\n",
            "      %8 : Float(125, 64, strides=[64, 1], requires_grad=0, device=cuda:0),\n",
            "      %9 : Float(125, 64, strides=[64, 1], requires_grad=0, device=cuda:0)):\n",
            "  %10 : Function = prim::Constant[name=\"batch_sw_func\"]()\n",
            "  %11 : Tensor = prim::CallFunction(%10, %0, %1, %2, %3, %4, %5, %6, %7, %8, %9)\n",
            "  return (%11)\n",
            "\n",
            "tensor([ 96.9170, 101.6837, 100.9182, 106.1258,  97.6396, 117.0418, 111.9109,\n",
            "         98.1991, 117.8617, 106.5082, 117.4011, 106.8881, 101.1058,  93.2988,\n",
            "        106.5921,  98.1852,  98.3521, 117.1855, 105.1182,  92.8305,  98.2466,\n",
            "        107.0196, 111.4845,  93.7431, 100.7904, 116.4669, 117.1199, 107.2120,\n",
            "        107.8658, 101.4239,  97.9903, 117.3321, 106.1214, 101.5301,  96.2617,\n",
            "        111.8204, 101.4959, 111.9469, 105.3003,  98.1917, 117.6322, 102.3818,\n",
            "         96.1251, 102.4135, 108.6407, 105.6136,  97.1293,  94.7116,  98.1466,\n",
            "         96.2407, 108.3316,  97.4162, 101.3835, 111.9653, 111.8093, 104.2527,\n",
            "         96.3414, 104.9545,  99.2517,  98.3273, 109.1088, 108.9274, 103.1509,\n",
            "        111.7842], device='cuda:0')\n",
            "avg torch.jit.trace batch: 0.06551995277404785 sec.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ko5_GHrdfIhY",
        "outputId": "d252d391-e045-4a4d-cd0f-30a2652919ba"
      },
      "source": [
        "assert torch_device == \"cuda\"\n",
        "\n",
        "# Warmup before capture  \n",
        "torch.cuda.synchronize()\n",
        "s = torch.cuda.Stream()\n",
        "s.wait_stream(torch.cuda.current_stream())\n",
        "with torch.cuda.stream(s):\n",
        "    for _ in range(3):\n",
        "        static_output = batch_sw(*example_inputs)\n",
        "torch.cuda.current_stream().wait_stream(s)\n",
        "torch.cuda.synchronize()\n",
        "\n",
        "# Captures the graph                                                                                                                           \n",
        "g = torch.cuda.CUDAGraph()\n",
        "with torch.cuda.graph(g):\n",
        "    static_output = batch_sw(*example_inputs)\n",
        "\n",
        "torch.cuda.synchronize()\n",
        "start = time.time()\n",
        "for i in range(n_trials):\n",
        "  g.replay()\n",
        "torch.cuda.synchronize()\n",
        "elapsed_time = time.time() - start\n",
        "print(static_output)\n",
        "print(\"avg torch.jit.trace and cuda graph batch :\", elapsed_time/n_trials, 'sec.')\n"
      ],
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([ 96.9170, 101.6837, 100.9182, 106.1258,  97.6396, 117.0418, 111.9109,\n",
            "         98.1991, 117.8617, 106.5082, 117.4011, 106.8881, 101.1058,  93.2988,\n",
            "        106.5921,  98.1852,  98.3521, 117.1855, 105.1182,  92.8305,  98.2466,\n",
            "        107.0196, 111.4845,  93.7431, 100.7904, 116.4669, 117.1199, 107.2120,\n",
            "        107.8658, 101.4239,  97.9903, 117.3321, 106.1214, 101.5301,  96.2617,\n",
            "        111.8204, 101.4959, 111.9469, 105.3003,  98.1917, 117.6322, 102.3818,\n",
            "         96.1251, 102.4135, 108.6407, 105.6136,  97.1293,  94.7116,  98.1466,\n",
            "         96.2407, 108.3316,  97.4162, 101.3835, 111.9653, 111.8093, 104.2527,\n",
            "         96.3414, 104.9545,  99.2517,  98.3273, 109.1088, 108.9274, 103.1509,\n",
            "        111.7842], device='cuda:0')\n",
            "avg torch.jit.trace and cuda graph batch : 0.031725740432739256 sec.\n"
          ]
        }
      ]
    }
  ]
}