HeadFirstLinearLayer.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "HeadFirstLinearLayer.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "s-O1wYxtRy0J",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import torch"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cFlYjo6zR44N",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "M = 128\n",
        "N = 128\n",
        "K = 128\n",
        "C = 32"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PIYODCpw3jtG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "matrix_a = np.random.randn(C, M, K)\n",
        "matrix_b = np.random.randn(C, K, N)\n",
        "matrix_c = np.random.randn(C, M, N)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G5rSww3bSw5h",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def matmul(m_a, m_b, m_c):\n",
        "  for c in range(C):\n",
        "    for m in range(M):\n",
        "      for n in range(N):\n",
        "        val = 0.\n",
        "        for k in range(K):\n",
        "          val += m_a[c,m,k] * m_b[c,k,n]\n",
        "        m_c[c,m,n] = val"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oXIB6Y0IT1ri",
        "colab_type": "code",
        "outputId": "02323780-c35c-49b9-a27b-8a5ff58f4289",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 50
        }
      },
      "source": [
        "%time matmul(matrix_a, matrix_b, matrix_c)"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CPU times: user 35.1 s, sys: 9.6 ms, total: 35.1 s\n",
            "Wall time: 35.1 s\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tMIeCd_eVFUv",
        "colab_type": "code",
        "outputId": "0f169f04-2c41-4f45-fec1-3acaa65da168",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "np.allclose(matrix_c, (matrix_a @ matrix_b))"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7vPQQuigT1wu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def matmul_boardcast(m_a, m_b, m_c):\n",
        "  for c in range(C):\n",
        "    for m in range(M):\n",
        "      for n in range(N):\n",
        "        m_c[c,m,n] = (m_a[c,m,:] * m_b[c,:,n]).sum()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZhK1JhYSSxvc",
        "colab_type": "code",
        "outputId": "0df23228-d58e-45fc-df62-9e41c9e582e4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 50
        }
      },
      "source": [
        "%time matmul_boardcast(matrix_a, matrix_b, matrix_c)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CPU times: user 2.14 s, sys: 78 ms, total: 2.22 s\n",
            "Wall time: 2.14 s\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B7G_5Em6Sxrd",
        "colab_type": "code",
        "outputId": "7e042149-4529-4e29-ccaf-2478d27b3f44",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "np.allclose(matrix_c, (matrix_a @ matrix_b))"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W4iY_09QWdBq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install pycuda"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZIYNqKsRWZ7b",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import pycuda.autoinit\n",
        "import pycuda.driver as drv\n",
        "from pycuda import gpuarray\n",
        "from pycuda.compiler import SourceModule"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "viNw3PpxWDjM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "kernel = SourceModule(\"\"\"\n",
        "__global__ void matmul(double *mat_a, double *mat_b, double *mat_c, int C, int M, int N, int K)\n",
        "{\n",
        "  int height = blockIdx.y * blockDim.y + threadIdx.y;\n",
        "  int weight = blockIdx.x * blockDim.x + threadIdx.x;\n",
        "  int channel = blockIdx.z * blockDim.z + threadIdx.z;\n",
        "  int thread_idx = channel * M * N + height * N + weight;\n",
        "\n",
        "  if (channel < C && height < M && weight < N) {\n",
        "    double val = 0;\n",
        "    for (int k = 0; k < K; k++)\n",
        "      val += mat_a[channel * M * N + height * N + k] * mat_b[channel * M * N + k * N + weight];\n",
        "    mat_c[thread_idx] = val;\n",
        "  }\n",
        "}\n",
        "\"\"\")\n",
        "\n",
        "def matmul_gpu(m_a, m_b, m_c):\n",
        "  dev_a = gpuarray.to_gpu(m_a.reshape(-1))\n",
        "  dev_b = gpuarray.to_gpu(m_b.reshape(-1))\n",
        "  dev_c = gpuarray.to_gpu(m_c.reshape(-1))\n",
        "  matmul_cuda = kernel.get_function(\"matmul\")\n",
        "  matmul_cuda(dev_a, dev_b, dev_c, np.int32(C), np.int32(M), np.int32(N), np.int32(K), block=(32,32,1), grid=(N//32,M//32,C))\n",
        "  return dev_c.get().reshape(C,M,N)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K2b5CbXIWENy",
        "colab_type": "code",
        "outputId": "76d5344a-c131-40b3-8858-cf06f9a8ea37",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 50
        }
      },
      "source": [
        "%time c = matmul_gpu(matrix_a, matrix_b, matrix_c)"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CPU times: user 7.18 ms, sys: 2 ms, total: 9.18 ms\n",
            "Wall time: 10.3 ms\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sVJQw5GLWEW7",
        "colab_type": "code",
        "outputId": "e0adbcba-2d0f-4e13-f815-e66ee2e9e2a2",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "np.allclose(c, (matrix_a @ matrix_b))"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fvVaFfiLHUx7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "matrix_a = torch.randn(C, M, K)\n",
        "matrix_b = torch.randn(C, K, N)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BjhfPktaJ6sV",
        "colab_type": "code",
        "outputId": "751d0c4a-9e47-496e-c431-3557f2aea6d6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 50
        }
      },
      "source": [
        "%time matrix_c = matrix_a.matmul(matrix_b)"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CPU times: user 3.38 ms, sys: 0 ns, total: 3.38 ms\n",
            "Wall time: 2.62 ms\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OrsFysYqKZTv",
        "colab_type": "code",
        "outputId": "0f5fc205-e50a-49a2-82b0-ad51fb47a8f0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "np.allclose(matrix_c, (matrix_a @ matrix_b))"
      ],
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 25
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "14EOi76yFZkl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}