Positive Matrix Product.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Positive matrix product"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Matrix-vector product vs Matrix element-wise multiplication by vector (cuda float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')\n",
    "n_in = 9216\n",
    "n_out = 128\n",
    "batch_size = 256\n",
    "W = torch.empty(n_out, n_in).normal_().to(device=device)\n",
    "x = torch.linspace(-10, 10, n_in).to(device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time: 0.70267ms\n"
     ]
    }
   ],
   "source": [
    "start = time.clock()\n",
    "result_1 = W @ x\n",
    "elapsed_time = time.clock() - start\n",
    "print(\"Elapsed time: \" + \"{0:0.5f}\".format(elapsed_time * 1000) + \"ms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time: 0.62806ms\n"
     ]
    }
   ],
   "source": [
    "start = time.clock()\n",
    "result_2 = W * x\n",
    "elapsed_time = time.clock() - start\n",
    "print(\"Elapsed time: \" + \"{0:0.5}\".format(elapsed_time * 1000) + \"ms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Matrix-vector product vs Matrix element-wise multiplication by vector (cpu int8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cpu')\n",
    "n_in = 9216\n",
    "n_out = 128\n",
    "batch_size = 256\n",
    "W = torch.empty(n_out, n_in, dtype=torch.int8).random_(-128, 127).to(device=device)\n",
    "x = torch.empty(n_in, dtype=torch.int8).random_(-128, 127).to(device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time: 1.07611ms\n"
     ]
    }
   ],
   "source": [
    "start = time.clock()\n",
    "result_1 = W @ x\n",
    "elapsed_time = time.clock() - start\n",
    "print(\"Elapsed time: \" + \"{0:0.5f}\".format(elapsed_time * 1000) + \"ms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time: 0.86689ms\n"
     ]
    }
   ],
   "source": [
    "start = time.clock()\n",
    "result_2 = (W * x)\n",
    "elapsed_time = time.clock() - start\n",
    "print(\"Elapsed time: \" + \"{0:0.5}\".format(elapsed_time * 1000) + \"ms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FC Layer VS manual matrix vector product VS element-wise multiplication + sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\")\n",
    "fc = nn.Linear(n_in, n_out, bias=False).to(device=device)\n",
    "for parameter in fc.parameters():\n",
    "    parameter.requires_grad = False  # This increases the speed of result 4 like CRAZY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = torch.Tensor(batch_size, n_in).normal_().to(device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time: 0.014s - Time per iteration: 0.141ms\n"
     ]
    }
   ],
   "source": [
    "n_iter = 100\n",
    "start = time.clock()\n",
    "for i in range(n_iter):\n",
    "    results = fc(inputs)\n",
    "elapsed_time = time.clock() - start\n",
    "print(\"Elapsed time: \" + \"{0:0.3f}\".format(elapsed_time) + \"s - Time per iteration: \" + \"{0:0.3f}\".format(elapsed_time/n_iter*1000) + \"ms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time: 0.067s - Time per iteration: 0.670ms\n"
     ]
    }
   ],
   "source": [
    "n_iter = 100\n",
    "W = fc.weight.data\n",
    "start = time.time()\n",
    "for i in range(n_iter):\n",
    "    result = W @ inputs.view(n_in, -1)\n",
    "    torch.cuda.synchronize()\n",
    "elapsed_time = time.time() - start\n",
    "print(\"Elapsed time: \" + \"{0:0.3f}\".format(elapsed_time) + \"s - Time per iteration: \" + \"{0:0.3f}\".format(elapsed_time/n_iter*1000) + \"ms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time: 3.675s - Time per iteration: 36.748ms\n"
     ]
    }
   ],
   "source": [
    "n_iter = 100\n",
    "W = fc.weight.data\n",
    "start = time.time()\n",
    "for i in range(n_iter):\n",
    "    inputs_expanded = inputs.unsqueeze(1).expand(batch_size, n_out, n_in)\n",
    "    result = (W * inputs_expanded).sum(axis=2)\n",
    "    torch.cuda.synchronize()\n",
    "elapsed_time = time.time() - start\n",
    "print(\"Elapsed time: \" + \"{0:0.3f}\".format(elapsed_time) + \"s - Time per iteration: \" + \"{0:0.3f}\".format(elapsed_time/n_iter*1000) + \"ms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computation of the signs matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\")\n",
    "W = torch.empty(n_out, n_in, dtype=torch.int8).random_(-128, 127).to(device)\n",
    "input_ = torch.empty(n_in, dtype=torch.int8).random_(-128, 127).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ -51,  112, -110,  ..., -103,   12,   -1],\n",
       "        [ -80,  119,   43,  ...,  119,  109,  -25],\n",
       "        [-124,  112,  -52,  ...,  -14,  -46,   10],\n",
       "        ...,\n",
       "        [ -12,  -29,   66,  ...,   90,  -92,   49],\n",
       "        [  -8,  -35,   58,  ...,  107,   83,  118],\n",
       "        [ 117, -100,   66,  ...,   42,   55,  111]], device='cuda:0',\n",
       "       dtype=torch.int8)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "W"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-118,  -36,   37,  ...,   97, -111,   71], device='cuda:0',\n",
       "       dtype=torch.int8)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1, -1, -1,  ..., -1, -1, -1],\n",
       "        [ 1, -1,  1,  ...,  1, -1, -1],\n",
       "        [ 1, -1, -1,  ..., -1,  1,  1],\n",
       "        ...,\n",
       "        [ 1,  1,  1,  ...,  1,  1,  1],\n",
       "        [ 1,  1,  1,  ...,  1, -1,  1],\n",
       "        [-1,  1,  1,  ...,  1, -1,  1]], device='cuda:0', dtype=torch.int8)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w_signs = torch.sign(W)  # The weight signs are precomputed (out of time benchmark)\n",
    "input_signs = torch.sign(input_)\n",
    "w_signs * input_signs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time: 0.326s - Time per iteration: 0.3262ms\n"
     ]
    }
   ],
   "source": [
    "n_iter = 1000\n",
    "start = time.time()\n",
    "if device == torch.device(\"cuda\"):\n",
    "    for i in range(n_iter):\n",
    "        input_signs = torch.sign(input_)  # A bit slower with cpu, a bit faster with cuda\n",
    "        signs_matrix = w_signs * input_signs\n",
    "else:\n",
    "    for i in range(n_iter):\n",
    "        input_signs = ((input_ >> 7) * -2 + 1)\n",
    "        signs_matrix = w_signs * input_signs\n",
    "\n",
    "elapsed_time = time.time() - start\n",
    "print(\"Elapsed time: \" + \"{0:0.3f}\".format(elapsed_time) + \"s - Time per iteration: \" + \"{0:0.4f}\".format(elapsed_time/n_iter * 1000) + \"ms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computation of the absolute value of the tensors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see the absolute value of the smallest number is bugged because it doesnt exist in that dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Absolute value of int8 -127 is 127\n",
      "Absolute value of int8 -128 is -128\n"
     ]
    }
   ],
   "source": [
    "value1 = torch.abs(torch.tensor(-127, dtype=torch.int8)).item()\n",
    "value2 = torch.abs(torch.tensor(-128, dtype=torch.int8)).item()\n",
    "print(\"Absolute value of int8 -127 is \" + repr(value1))\n",
    "print(\"Absolute value of int8 -128 is \" + repr(value2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cpu\")\n",
    "W = torch.empty(n_out, n_in, dtype=torch.int8).random_(-127, 127).to(device)  # We only range from -127 to avoid bug of abs\n",
    "input_ = torch.empty(n_in, dtype=torch.int8).random_(-127, 127).to(device)  # Same"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "abs_W = torch.abs(W)  # This can be precomputed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Elapsed time: 0.188s - Time per iteration: 0.0188ms\n"
     ]
    }
   ],
   "source": [
    "n_iter = 10000\n",
    "start = time.time()\n",
    "for i in range(n_iter):\n",
    "    torch.abs(input_)\n",
    "elapsed_time = time.time() - start\n",
    "print(\"Elapsed time: \" + \"{0:0.3f}\".format(elapsed_time) + \"s - Time per iteration: \" + \"{0:0.4f}\".format(elapsed_time/n_iter*1000) + \"ms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Full pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 256\n",
    "n_in = 9216\n",
    "n_out = 128\n",
    "device = torch.device(\"cpu\")\n",
    "\n",
    "# Generate W and precompute it's element-wise signs and absolute values\n",
    "W = torch.empty(n_out, n_in, dtype=torch.int8, device=device).random_(-127, 127)\n",
    "abs_W = torch.abs(W)\n",
    "signs_W = torch.sign(W)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate batch_size inputs\n",
    "inputs = torch.empty(batch_size, n_in, dtype=torch.int8, device=device).random_(-127, 127)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "abs_inputs_expanded = torch.abs(inputs).unsqueeze(1).expand(batch_size, n_out, n_in)\n",
    "signs_inputs_expanded = torch.sign(inputs).unsqueeze(1).expand(batch_size, n_out, n_in)\n",
    "signs_matrices = signs_W * signs_inputs_expanded\n",
    "\n",
    "# The first multiplication (abs_W * abs_inputs_expanded) can use only fib encoded numbers\n",
    "# The second multiplication (... * signs_matrices) is always an unsigned int * a sign\n",
    "result = torch.zeros_like(abs_W, dtype=torch.int32)\n",
    "torch.mul(abs_W, abs_inputs_expanded, out=result)\n",
    "torch.mul(result, signs_matrices, out=result)\n",
    "result = result.sum(dim=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[   55689,  -225734,  -195977,  ...,   680322,  -100576,   571251],\n",
       "        [  546280,   455686,  -596431,  ...,  -187095,  -641463,   722258],\n",
       "        [ -177177,  -105483,  -774910,  ...,  -303105,   721467,  -356380],\n",
       "        ...,\n",
       "        [ -648762, -1156661,   570426,  ...,  -660435,   658413,  1071726],\n",
       "        [  399636,  -708527, -1157746,  ...,  -721120,  -101464,   495962],\n",
       "        [     303,   474557,   201979,  ...,  -708200,   173130,  -564593]])"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\")\n",
    "inputs_cuda = inputs.to(dtype=torch.float32, device=device)\n",
    "fc = nn.Linear(n_in, n_out, bias=False).to(device=device)\n",
    "with torch.no_grad():\n",
    "    fc.weight.data = W.to(dtype=torch.float32, device=device)\n",
    "    result_baseline = fc(inputs_cuda)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 5.5689e+04, -2.2573e+05, -1.9598e+05,  ...,  6.8032e+05,\n",
       "         -1.0058e+05,  5.7125e+05],\n",
       "        [ 5.4628e+05,  4.5569e+05, -5.9643e+05,  ..., -1.8710e+05,\n",
       "         -6.4146e+05,  7.2226e+05],\n",
       "        [-1.7718e+05, -1.0548e+05, -7.7491e+05,  ..., -3.0310e+05,\n",
       "          7.2147e+05, -3.5638e+05],\n",
       "        ...,\n",
       "        [-6.4876e+05, -1.1567e+06,  5.7043e+05,  ..., -6.6044e+05,\n",
       "          6.5841e+05,  1.0717e+06],\n",
       "        [ 3.9964e+05, -7.0853e+05, -1.1577e+06,  ..., -7.2112e+05,\n",
       "         -1.0146e+05,  4.9596e+05],\n",
       "        [ 3.0300e+02,  4.7456e+05,  2.0198e+05,  ..., -7.0820e+05,\n",
       "          1.7313e+05, -5.6459e+05]], device='cuda:0')"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result_baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check equality between baseline and positive matrix product trick\n",
    "torch.all(torch.eq(result.to(dtype=torch.float32, device=device), result_baseline)).item()"
   ]
  }
 ],
 "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.6.7"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}