{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hwj/.miniconda3/envs/pyg/lib/python3.8/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "from torch import Tensor\n",
    "from typing import *\n",
    "\n",
    "from torch_scatter import scatter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def act_blks(act_mask: Tensor, blk_size: int) -> Tensor:\n",
    "    assert act_mask.dtype == torch.bool\n",
    "    assert act_mask.dim() == 1\n",
    "    n = act_mask.size(0)\n",
    "    m = (n + blk_size - 1) // blk_size * blk_size\n",
    "\n",
    "    blk_mask = act_mask.detach()\n",
    "    if n != m:\n",
    "        blk_mask = torch.zeros(m, dtype=act_mask.dtype, device=act_mask.device)\n",
    "        blk_mask[:n] = act_mask.detach()\n",
    "    return blk_mask.view(-1, blk_size).max(dim=-1).values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "act_mask = torch.rand(1000) < 0.01\n",
    "blk_mask = act_blks(act_mask, 64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "cached_blocks = torch.randn(blk_mask.size(0), 64, 8)\n",
    "blks = cached_blocks[blk_mask]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gather_block(blks: Tensor, blk_mask: Tensor, index: Tensor) -> Tensor:\n",
    "    blk_size = blks.size(1)\n",
    "    blk_mask = blk_mask.type(torch.long).cumsum_(dim=0) * blk_mask - 1\n",
    "    \n",
    "    idx0 = index.div(blk_size, rounding_mode=\"floor\")\n",
    "    ind_mask = blk_mask[idx0] >= 0\n",
    "    \n",
    "    idx0 = blk_mask[idx0][ind_mask]\n",
    "    idx1 = index[ind_mask] % blk_size\n",
    "\n",
    "    s = index.shape[:1] + blks.shape[2:]\n",
    "    data = torch.zeros(s, dtype=blks.dtype, device=blks.device)\n",
    "    data[ind_mask] = blks[idx0, idx1, ...]\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scatter_block(x: Tensor, index: Tensor, blk_mask: Tensor, blk_size: int, reduce: str = \"sum\") -> Tensor:\n",
    "    num_blks = blk_mask.count_nonzero()\n",
    "    blk_mask = blk_mask.type(torch.long).cumsum_(dim=0) * blk_mask - 1\n",
    "\n",
    "    idx0 = index.div(blk_size, rounding_mode=\"floor\")\n",
    "    ind_mask = blk_mask[idx0] >= 0\n",
    "\n",
    "    idx0 = blk_mask[idx0][ind_mask]\n",
    "    idx1 = index[ind_mask] % blk_size\n",
    "\n",
    "    data = scatter(\n",
    "        src=x[ind_mask],\n",
    "        index=idx0*blk_size+idx1,\n",
    "        dim=0,\n",
    "        dim_size=num_blks*blk_size,\n",
    "        reduce=reduce,\n",
    "    )\n",
    "    return data.view(num_blks, blk_size, *x.shape[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def index_block(index: Tensor, blk_mask: Tensor, blk_size: int) -> Tensor:\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([False,  True,  True, False,  True, False, False,  True, False,  True,\n",
       "        False,  True,  True, False, False, False, False, False, False, False,\n",
       "        False,  True,  True, False, False, False,  True, False, False,  True,\n",
       "         True, False, False,  True,  True, False,  True,  True, False, False,\n",
       "         True, False,  True, False, False, False, False, False, False, False,\n",
       "         True,  True, False, False, False,  True,  True,  True, False, False,\n",
       "         True, False,  True, False,  True, False,  True, False, False,  True,\n",
       "        False,  True,  True,  True, False, False, False, False, False,  True,\n",
       "        False,  True, False,  True,  True,  True, False,  True, False,  True,\n",
       "        False,  True, False,  True,  True, False, False,  True, False,  True])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index = torch.randint(1000, size=(100,), dtype=torch.long)\n",
    "(gather_block(blks, blk_mask, index)**2).sum(dim=-1) != 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 64, 8])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(100, 8)\n",
    "scatter_block(x, index, blk_mask, 64).size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([True, True, True, True, True, True, True, True])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t = scatter_block(x, index, blk_mask, 64)\n",
    "t = t.view(-1, *t.shape[2:])\n",
    "act_blks((t**2).sum(dim=-1) != 0, 64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.return_types.sort(\n",
       "values=tensor([  5,  15,  18,  19,  25,  48,  62,  64,  68,  76,  88,  92,  94, 106,\n",
       "        114, 115, 120, 122, 126, 131, 135, 165, 169, 169, 177, 192, 214, 215,\n",
       "        245, 249, 255, 269, 282, 286, 294, 299, 334, 358, 359, 372, 401, 403,\n",
       "        430, 437, 455, 459, 459, 459, 471, 484, 490, 504, 505, 516, 528, 529,\n",
       "        559, 584, 613, 614, 614, 624, 629, 637, 638, 674, 679, 701, 706, 735,\n",
       "        742, 746, 748, 750, 759, 774, 774, 776, 778, 787, 803, 825, 831, 834,\n",
       "        861, 862, 868, 877, 882, 889, 899, 909, 921, 933, 938, 954, 960, 969,\n",
       "        981, 986]),\n",
       "indices=tensor([ 6, 77, 14, 41, 43,  0, 13, 53, 38, 54, 15, 76, 98, 47, 75, 80,  8, 95,\n",
       "        17, 34, 83, 87, 62, 93, 81, 20, 67, 39, 32,  5, 70,  4, 89, 72, 37, 64,\n",
       "        94, 42, 57, 29, 56, 21, 60,  7, 45, 59, 86, 96, 49, 92, 28, 88, 24, 68,\n",
       "         3, 61, 44,  2, 79, 55, 85, 33, 12,  1, 11, 97, 50,  9, 84, 40, 71, 22,\n",
       "        26, 30, 73, 16, 90, 27, 19, 35, 48, 18, 46, 74, 65, 78, 52, 31, 23, 58,\n",
       "        91, 66, 99, 69, 51, 36, 82, 63, 25, 10]))"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index.sort()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mpi",
   "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.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
