GEMM.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pyopencl as cl\n",
    "import os\n",
    "import time\n",
    "\n",
    "os.environ['PYOPENCL_COMPILER_OUTPUT'] = '0'\n",
    "os.environ['PYOPENCL_CTX'] = '0:1'\n",
    "\n",
    "class Timer:\n",
    "    def __enter__(self):\n",
    "        self.start = time.time()\n",
    "        return self\n",
    "\n",
    "    def __exit__(self, *args):\n",
    "        self.end = time.time()\n",
    "        self.interval = self.end - self.start\n",
    "        \n",
    "def roundUp(numToRound, multiple):\n",
    "    if (multiple == 0):\n",
    "        return numToRound\n",
    "    remainder = numToRound % multiple\n",
    "    if (remainder == 0):\n",
    "        return numToRound\n",
    "    return numToRound + multiple - remainder\n",
    "\n",
    "BLOCK_SIZE = 16\n",
    "WPT = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "kernel = \"\"\"\n",
    "#define BLOCK_SIZE 16\n",
    "#define global_idx(x_idx, y_idx, m) (x_idx * m + y_idx)\n",
    "\n",
    "#define WPT 8\n",
    "#define RBLOCK_SIZE (BLOCK_SIZE/WPT)\n",
    "\n",
    "// GEMM -- \n",
    "// A is M x N\n",
    "// B is N x P\n",
    "// C (output) is M x P\n",
    "\n",
    "// -- Uses NDRange Kernel with Local Memory\n",
    "// M, N, P must all be multiples of BLOCKSIZE\n",
    "__kernel __attribute__((reqd_work_group_size(BLOCK_SIZE, BLOCK_SIZE, 1)))\n",
    "void GEMM(\n",
    "      __global float* restrict A,\n",
    "      __global float* restrict B,\n",
    "      __global float* restrict C,\n",
    "      __const int M,\n",
    "      __const int N,\n",
    "      __const int P)\n",
    "{\n",
    "    const int row = get_local_id(0);\n",
    "    const int col = get_local_id(1);\n",
    "    const int m = BLOCK_SIZE*get_group_id(0) + row;\n",
    "    const int p = BLOCK_SIZE*get_group_id(1) + col;    \n",
    "    __local float A_local[BLOCK_SIZE][BLOCK_SIZE];\n",
    "    __local float B_local[BLOCK_SIZE][BLOCK_SIZE];\n",
    "    float acc = 0.0f;\n",
    "    const int numTiles = N/BLOCK_SIZE;\n",
    "    #pragma unroll\n",
    "    for (int t=0; t<numTiles; t++) {\n",
    "        const int r = BLOCK_SIZE*t + row;\n",
    "        const int c = BLOCK_SIZE*t + col;\n",
    "        A_local[row][col] = A[m*N + c];\n",
    "        B_local[row][col] = B[r*P + p];\n",
    "        barrier(CLK_LOCAL_MEM_FENCE);\n",
    "        #pragma unroll BLOCK_SIZE\n",
    "        for (int k=0; k<BLOCK_SIZE; k++){{\n",
    "            acc += A_local[row][k] * B_local[k][col];\n",
    "        }}\n",
    "        barrier(CLK_LOCAL_MEM_FENCE);\n",
    "    }\n",
    "    C[m*P + p] = acc;\n",
    "}\n",
    "\n",
    "// -- Uses NDRange Kernel with Local Memory and 1D Register tiling\n",
    "// M, N, P must all be multiples of BLOCKSIZE\n",
    "__kernel __attribute__((reqd_work_group_size(RBLOCK_SIZE, BLOCK_SIZE, 1)))\n",
    "void GEMM_1DREG(\n",
    "      __global float* restrict A,\n",
    "      __global float* restrict B,\n",
    "      __global float* restrict C,\n",
    "      __const int M,\n",
    "      __const int N,\n",
    "      __const int P)\n",
    "{\n",
    "    const int row = get_local_id(0);\n",
    "    const int col = get_local_id(1);\n",
    "    const int m = BLOCK_SIZE*get_group_id(0) + row;\n",
    "    const int p = BLOCK_SIZE*get_group_id(1) + col;    \n",
    "    __local float A_local[BLOCK_SIZE][BLOCK_SIZE];\n",
    "    __local float B_local[BLOCK_SIZE][BLOCK_SIZE];\n",
    "    \n",
    "    float acc[WPT];\n",
    "    for(int w=0; w<WPT; w++){\n",
    "        acc[w] = 0.0f;\n",
    "    }\n",
    "    const int numTiles = N/BLOCK_SIZE;\n",
    "    #pragma unroll\n",
    "    for (int t=0; t<numTiles; t++) {\n",
    "        for (int w=0; w<WPT; w++){\n",
    "            const int r = BLOCK_SIZE*t + row;\n",
    "            const int c = BLOCK_SIZE*t + col;\n",
    "            A_local[row + w*RBLOCK_SIZE][col] = A[(m + w*RBLOCK_SIZE)*N + c];\n",
    "            B_local[row + w*RBLOCK_SIZE][col] = B[(r + w*RBLOCK_SIZE)*P + p];\n",
    "        }\n",
    "        barrier(CLK_LOCAL_MEM_FENCE);\n",
    "        #pragma unroll BLOCK_SIZE\n",
    "        for (int k=0; k<BLOCK_SIZE; k++){\n",
    "            for (int w=0; w<WPT; w++){\n",
    "                acc[w] += A_local[row + w*RBLOCK_SIZE][k] * B_local[k][col];\n",
    "            }\n",
    "        }\n",
    "        barrier(CLK_LOCAL_MEM_FENCE);\n",
    "    }\n",
    "    for (int w=0; w<WPT; w++){\n",
    "        C[(m + w*RBLOCK_SIZE)*P + p] = acc[w];\n",
    "    }\n",
    "}\n",
    "\n",
    "// -- Uses NDRange Kernel with Local Memory and 2D Register tiling\n",
    "__kernel __attribute__((reqd_work_group_size(BLOCK_SIZE / WPT, BLOCK_SIZE / WPT, 1)))\n",
    "void GEMM_2DREG(\n",
    "      __global float* restrict A,\n",
    "      __global float* restrict B,\n",
    "      __global float* restrict C,\n",
    "      __const int M,\n",
    "      __const int N,\n",
    "      __const int P)\n",
    "{\n",
    "    const int row = get_local_id(0);\n",
    "    const int col = get_local_id(1);\n",
    "    const int m = BLOCK_SIZE*get_group_id(0) + row;\n",
    "    const int p = BLOCK_SIZE*get_group_id(1) + col;\n",
    "    __local float A_local[BLOCK_SIZE][BLOCK_SIZE];\n",
    "    __local float B_local[BLOCK_SIZE][BLOCK_SIZE];\n",
    "    \n",
    "    float Areg;\n",
    "    float Breg[WPT];\n",
    "    float acc[WPT][WPT];\n",
    "    for(int wm=0; wm<WPT; wm++){\n",
    "        for(int wn=0; wn<WPT; wn++){\n",
    "            acc[wm][wn] = 0.0f;\n",
    "        }\n",
    "    }\n",
    "    const int numTiles = N/BLOCK_SIZE;\n",
    "    #pragma unroll\n",
    "    for (int t=0; t<numTiles; t++) {\n",
    "        for (int wm=0; wm<WPT; wm++){\n",
    "            for (int wn=0; wn<WPT; wn++){\n",
    "                const int r = BLOCK_SIZE*t + row;\n",
    "                const int c = BLOCK_SIZE*t + col;\n",
    "                A_local[row + wm*RBLOCK_SIZE][col + wn*RBLOCK_SIZE] = A[(m + wm*RBLOCK_SIZE)*N + (c + wn*RBLOCK_SIZE)];\n",
    "                B_local[row + wm*RBLOCK_SIZE][col + wn*RBLOCK_SIZE] = B[(r + wm*RBLOCK_SIZE)*P + (p + wn*RBLOCK_SIZE)];\n",
    "            }\n",
    "        }\n",
    "        barrier(CLK_LOCAL_MEM_FENCE);\n",
    "        #pragma unroll BLOCK_SIZE\n",
    "        for (int k=0; k<BLOCK_SIZE; k++){\n",
    "            for (int wn=0; wn<WPT; wn++){\n",
    "                Breg[wn] = B_local[k][col + wn*RBLOCK_SIZE];\n",
    "            }\n",
    "            for (int wm=0; wm<WPT; wm++){\n",
    "                Areg = A_local[row + wm*RBLOCK_SIZE][k];\n",
    "                for (int wn=0; wn<WPT; wn++){\n",
    "                    acc[wm][wn] += Areg * Breg[wn];\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "        barrier(CLK_LOCAL_MEM_FENCE);\n",
    "    }\n",
    "    for (int wm=0; wm<WPT; wm++){\n",
    "        for (int wn=0; wn<WPT; wn++){\n",
    "            C[(m + wm*RBLOCK_SIZE)*P + (p + wn*RBLOCK_SIZE)] = acc[wm][wn];\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "// -- Uses NDRange Kernel with Local Memory and 2D Register tiling\n",
    "// M, N, P can be arbitrary sizes\n",
    "__kernel __attribute__((reqd_work_group_size(BLOCK_SIZE / WPT, BLOCK_SIZE / WPT, 1)))\n",
    "void GEMM_2DREG_IMITATE_PADDING(\n",
    "      __global float* restrict A,\n",
    "      __global float* restrict B,\n",
    "      __global float* restrict C,\n",
    "      __const int M,\n",
    "      __const int N,\n",
    "      __const int P,\n",
    "      __const int M_,\n",
    "      __const int N_,\n",
    "      __const int P_)\n",
    "{\n",
    "    const int row = get_local_id(0);\n",
    "    const int col = get_local_id(1);\n",
    "    const int m = BLOCK_SIZE*get_group_id(0) + row;\n",
    "    const int p = BLOCK_SIZE*get_group_id(1) + col;\n",
    "    __local float A_local[BLOCK_SIZE][BLOCK_SIZE];\n",
    "    __local float B_local[BLOCK_SIZE][BLOCK_SIZE];\n",
    "    \n",
    "    float Areg;\n",
    "    float Breg[WPT];\n",
    "    float acc[WPT][WPT];\n",
    "    for(int wm=0; wm<WPT; wm++){\n",
    "        for(int wn=0; wn<WPT; wn++){\n",
    "            acc[wm][wn] = 0.0f;\n",
    "        }\n",
    "    }\n",
    "    const int numTiles = N_/BLOCK_SIZE;\n",
    "    #pragma unroll\n",
    "    for (int t=0; t<numTiles; t++) {\n",
    "        for (int wm=0; wm<WPT; wm++){\n",
    "            for (int wn=0; wn<WPT; wn++){\n",
    "                const int r = BLOCK_SIZE*t + row;\n",
    "                const int c = BLOCK_SIZE*t + col;\n",
    "                if(((m + wm*RBLOCK_SIZE) < M) && ((c + wn*RBLOCK_SIZE) < N)){\n",
    "                    A_local[row + wm*RBLOCK_SIZE][col + wn*RBLOCK_SIZE] = A[(m + wm*RBLOCK_SIZE)*N + (c + wn*RBLOCK_SIZE)];\n",
    "                } else {\n",
    "                    A_local[row + wm*RBLOCK_SIZE][col + wn*RBLOCK_SIZE] = 0.0;\n",
    "                }\n",
    "\n",
    "                if(((p + wn*RBLOCK_SIZE) < P) && ((r + wm*RBLOCK_SIZE) < N)){\n",
    "                    B_local[row + wm*RBLOCK_SIZE][col + wn*RBLOCK_SIZE] = B[(r + wm*RBLOCK_SIZE)*P + (p + wn*RBLOCK_SIZE)];\n",
    "                } else {\n",
    "                    B_local[row + wm*RBLOCK_SIZE][col + wn*RBLOCK_SIZE] = 0.0;\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "        barrier(CLK_LOCAL_MEM_FENCE);\n",
    "        #pragma unroll BLOCK_SIZE\n",
    "        for (int k=0; k<BLOCK_SIZE; k++){\n",
    "            for (int wn=0; wn<WPT; wn++){\n",
    "                Breg[wn] = B_local[k][col + wn*RBLOCK_SIZE];\n",
    "            }\n",
    "            for (int wm=0; wm<WPT; wm++){\n",
    "                Areg = A_local[row + wm*RBLOCK_SIZE][k];\n",
    "                for (int wn=0; wn<WPT; wn++){\n",
    "                    acc[wm][wn] += Areg * Breg[wn];\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "        barrier(CLK_LOCAL_MEM_FENCE);\n",
    "    }\n",
    "    for (int wm=0; wm<WPT; wm++){\n",
    "        for (int wn=0; wn<WPT; wn++){\n",
    "            if(((m + wm*RBLOCK_SIZE) < M) && ((p + wn*RBLOCK_SIZE) < P)){\n",
    "                C[(m + wm*RBLOCK_SIZE)*P + (p + wn*RBLOCK_SIZE)] = acc[wm][wn];\n",
    "            }\n",
    "            \n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "// -- Uses NDRange Kernel with Local Memory\n",
    "// M, N, P can be arbitrary sizes\n",
    "__kernel __attribute__((reqd_work_group_size(BLOCK_SIZE, BLOCK_SIZE, 1)))\n",
    "void GEMM_IMITATE_PADDING(\n",
    "      __global float* restrict A,\n",
    "      __global float* restrict B,\n",
    "      __global float* restrict C,\n",
    "      __const int M,\n",
    "      __const int N,\n",
    "      __const int P,\n",
    "      __const int M_,\n",
    "      __const int N_,\n",
    "      __const int P_)\n",
    "{\n",
    "    const int row = get_local_id(0);\n",
    "    const int col = get_local_id(1);\n",
    "    const int m = BLOCK_SIZE*get_group_id(0) + row;\n",
    "    const int p = BLOCK_SIZE*get_group_id(1) + col;\n",
    "     \n",
    "    __local float A_local[BLOCK_SIZE][BLOCK_SIZE];\n",
    "    __local float B_local[BLOCK_SIZE][BLOCK_SIZE];\n",
    "    float acc = 0.0f;\n",
    "    const int numTiles = N_/BLOCK_SIZE;\n",
    "    \n",
    "    #pragma unroll 4\n",
    "    \n",
    "    for (int t=0; t<numTiles; t++) {\n",
    "        const int r = BLOCK_SIZE*t + row;\n",
    "        const int c = BLOCK_SIZE*t + col;\n",
    "        \n",
    "        if((m < M) && (c < N)){\n",
    "            int A_idx = m*N + c;\n",
    "            A_local[row][col] = A[A_idx];\n",
    "        } else {\n",
    "            A_local[row][col] = 0.0;\n",
    "        }\n",
    "        \n",
    "        if((p < P) && (r < N)){\n",
    "            int B_idx = r*P + p;\n",
    "            B_local[row][col] = B[B_idx];\n",
    "        } else {\n",
    "            B_local[row][col] = 0.0;\n",
    "        }\n",
    "        \n",
    "        barrier(CLK_LOCAL_MEM_FENCE);\n",
    "        #pragma unroll BLOCK_SIZE\n",
    "        for (int k=0; k<BLOCK_SIZE; k++){{\n",
    "            acc += A_local[row][k] * B_local[k][col];\n",
    "        }}\n",
    "        barrier(CLK_LOCAL_MEM_FENCE);\n",
    "    }\n",
    "    \n",
    "    if((m < M) && (p < P)){\n",
    "        C[m*P + p] = acc;\n",
    "    }\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MATRIX MULTIPLICATION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "np:  0.0015859603881835938\n",
      "py:  0.0009908676147460938\n",
      "cl:  0.06376481056213379\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "((304, 512), 304, 400, 512)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mult = 2\n",
    "M, N, P = 3*(10**mult), 4*(10**mult), 5*(10**mult)\n",
    "M, N, P = roundUp(M, BLOCK_SIZE), roundUp(N, BLOCK_SIZE), roundUp(P, BLOCK_SIZE)\n",
    "A = np.random.randn(M, N).astype('float32')\n",
    "B = np.random.randn(N, P).astype('float32')\n",
    "ctx = cl.create_some_context()\n",
    "queue = cl.CommandQueue(ctx)\n",
    "mf = cl.mem_flags\n",
    "a_device = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=A)\n",
    "b_device = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=B)\n",
    "\n",
    "prg = cl.Program(ctx, kernel).build()\n",
    "start = time.time()\n",
    "C = A.dot(B)\n",
    "print('np: ', time.time() - start)\n",
    "start = time.time()\n",
    "C = A @ B\n",
    "print('py: ', time.time() - start)\n",
    "\n",
    "C_out = np.zeros_like(C)\n",
    "c_device = cl.Buffer(ctx, mf.WRITE_ONLY, C_out.nbytes)\n",
    "C.shape\n",
    "\n",
    "with Timer() as t:\n",
    "    event = prg.GEMM(queue, (M, P), (BLOCK_SIZE, BLOCK_SIZE), \n",
    "                     a_device, b_device, c_device,\n",
    "                     np.int32(M), np.int32(N), np.int32(P))\n",
    "    event.wait()\n",
    "print('cl: ', t.interval)\n",
    "\n",
    "cl.enqueue_copy(queue, C_out, c_device)\n",
    "C_out.shape, M, N, P"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  9.439566  ,  12.805741  , -10.724502  , ..., -22.892948  ,\n",
       "          6.8262587 ,  12.482617  ],\n",
       "       [ -4.716728  ,   3.5809178 ,   8.460974  , ..., -11.874376  ,\n",
       "        -14.738298  , -19.674625  ],\n",
       "       [ 10.3869    ,   0.33998686,  -2.2919376 , ...,  -3.548169  ,\n",
       "         -4.678543  ,   1.2437406 ],\n",
       "       ...,\n",
       "       [  5.6830344 ,  14.351683  ,  46.228428  , ...,  -4.9663544 ,\n",
       "         13.358488  ,   6.82675   ],\n",
       "       [ 25.55392   ,  -7.1336083 ,  10.141382  , ...,   9.633443  ,\n",
       "        -25.31231   ,  -2.9148455 ],\n",
       "       [ 24.146011  ,   6.7507944 ,  18.584644  , ...,  14.319832  ,\n",
       "          3.0579646 ,  29.500877  ]], dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  9.439571  ,  12.805742  , -10.724497  , ..., -22.892923  ,\n",
       "          6.8262577 ,  12.482628  ],\n",
       "       [ -4.716733  ,   3.5809164 ,   8.46097   , ..., -11.874374  ,\n",
       "        -14.738298  , -19.674622  ],\n",
       "       [ 10.386899  ,   0.33999014,  -2.2919402 , ...,  -3.5481699 ,\n",
       "         -4.678542  ,   1.2437413 ],\n",
       "       ...,\n",
       "       [  5.683032  ,  14.351689  ,  46.228428  , ...,  -4.9663506 ,\n",
       "         13.358485  ,   6.826749  ],\n",
       "       [ 25.553919  ,  -7.1336064 ,  10.141384  , ...,   9.63344   ,\n",
       "        -25.312307  ,  -2.9148445 ],\n",
       "       [ 24.146008  ,   6.7507973 ,  18.584654  , ...,  14.31982   ,\n",
       "          3.0579624 ,  29.500868  ]], dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MATRIX MULTIPLICATION 1D REGISTER TILING"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "np:  0.0018579959869384766\n",
      "py:  0.0016303062438964844\n",
      "cl:  0.08417892456054688\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "((304, 512), 304, 400, 512)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ctx = cl.create_some_context()\n",
    "queue = cl.CommandQueue(ctx)\n",
    "mf = cl.mem_flags\n",
    "a_device = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=A)\n",
    "b_device = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=B)\n",
    "\n",
    "prg = cl.Program(ctx, kernel).build()\n",
    "start = time.time()\n",
    "C = A.dot(B)\n",
    "print('np: ', time.time() - start)\n",
    "start = time.time()\n",
    "C = A @ B\n",
    "print('py: ', time.time() - start)\n",
    "\n",
    "C_out = np.zeros_like(C)\n",
    "c_device = cl.Buffer(ctx, mf.WRITE_ONLY, C_out.nbytes)\n",
    "C.shape\n",
    "\n",
    "with Timer() as t:\n",
    "    event = prg.GEMM_1DREG(queue, (M // WPT, P), (BLOCK_SIZE // WPT, BLOCK_SIZE), \n",
    "                     a_device, b_device, c_device,\n",
    "                     np.int32(M), np.int32(N), np.int32(P))\n",
    "    event.wait()\n",
    "print('cl: ', t.interval)\n",
    "\n",
    "cl.enqueue_copy(queue, C_out, c_device)\n",
    "C_out.shape, M, N, P"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  9.439571  ,  12.805742  , -10.724497  , ..., -22.892923  ,\n",
       "          6.8262577 ,  12.482628  ],\n",
       "       [ -4.716733  ,   3.5809164 ,   8.46097   , ..., -11.874374  ,\n",
       "        -14.738298  , -19.674622  ],\n",
       "       [ 10.386899  ,   0.33999014,  -2.2919402 , ...,  -3.5481699 ,\n",
       "         -4.678542  ,   1.2437413 ],\n",
       "       ...,\n",
       "       [  5.683032  ,  14.351689  ,  46.228428  , ...,  -4.9663506 ,\n",
       "         13.358485  ,   6.826749  ],\n",
       "       [ 25.553919  ,  -7.1336064 ,  10.141384  , ...,   9.63344   ,\n",
       "        -25.312307  ,  -2.9148445 ],\n",
       "       [ 24.146008  ,   6.7507973 ,  18.584654  , ...,  14.31982   ,\n",
       "          3.0579624 ,  29.500868  ]], dtype=float32)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  9.439566  ,  12.805741  , -10.724502  , ..., -22.892948  ,\n",
       "          6.8262587 ,  12.482617  ],\n",
       "       [ -4.716728  ,   3.5809178 ,   8.460974  , ..., -11.874376  ,\n",
       "        -14.738298  , -19.674625  ],\n",
       "       [ 10.3869    ,   0.33998686,  -2.2919376 , ...,  -3.548169  ,\n",
       "         -4.678543  ,   1.2437406 ],\n",
       "       ...,\n",
       "       [  5.6830344 ,  14.351683  ,  46.228428  , ...,  -4.9663544 ,\n",
       "         13.358488  ,   6.82675   ],\n",
       "       [ 25.55392   ,  -7.1336083 ,  10.141382  , ...,   9.633443  ,\n",
       "        -25.31231   ,  -2.9148455 ],\n",
       "       [ 24.146011  ,   6.7507944 ,  18.584644  , ...,  14.319832  ,\n",
       "          3.0579646 ,  29.500877  ]], dtype=float32)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MATRIX MULTIPLICATION REGISTER TILING"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "np:  0.0019860267639160156\n",
      "py:  0.0014548301696777344\n",
      "cl:  0.13943099975585938\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "((304, 512), 304, 400, 512)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ctx = cl.create_some_context()\n",
    "queue = cl.CommandQueue(ctx)\n",
    "mf = cl.mem_flags\n",
    "a_device = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=A)\n",
    "b_device = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=B)\n",
    "\n",
    "prg = cl.Program(ctx, kernel).build()\n",
    "start = time.time()\n",
    "C = A.dot(B)\n",
    "print('np: ', time.time() - start)\n",
    "start = time.time()\n",
    "C = A @ B\n",
    "print('py: ', time.time() - start)\n",
    "\n",
    "C_out = np.zeros_like(C)\n",
    "c_device = cl.Buffer(ctx, mf.WRITE_ONLY, C_out.nbytes)\n",
    "C.shape\n",
    "\n",
    "with Timer() as t:\n",
    "    event = prg.GEMM_2DREG(queue, (M // WPT, P // WPT), (BLOCK_SIZE // WPT, BLOCK_SIZE // WPT), \n",
    "                     a_device, b_device, c_device,\n",
    "                     np.int32(M), np.int32(N), np.int32(P))\n",
    "    event.wait()\n",
    "print('cl: ', t.interval)\n",
    "\n",
    "cl.enqueue_copy(queue, C_out, c_device)\n",
    "C_out.shape, M, N, P"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  9.439571  ,  12.805742  , -10.724497  , ..., -22.892923  ,\n",
       "          6.8262577 ,  12.482628  ],\n",
       "       [ -4.716733  ,   3.5809164 ,   8.46097   , ..., -11.874374  ,\n",
       "        -14.738298  , -19.674622  ],\n",
       "       [ 10.386899  ,   0.33999014,  -2.2919402 , ...,  -3.5481699 ,\n",
       "         -4.678542  ,   1.2437413 ],\n",
       "       ...,\n",
       "       [  5.683032  ,  14.351689  ,  46.228428  , ...,  -4.9663506 ,\n",
       "         13.358485  ,   6.826749  ],\n",
       "       [ 25.553919  ,  -7.1336064 ,  10.141384  , ...,   9.63344   ,\n",
       "        -25.312307  ,  -2.9148445 ],\n",
       "       [ 24.146008  ,   6.7507973 ,  18.584654  , ...,  14.31982   ,\n",
       "          3.0579624 ,  29.500868  ]], dtype=float32)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  9.439566  ,  12.805741  , -10.724502  , ..., -22.892948  ,\n",
       "          6.8262587 ,  12.482617  ],\n",
       "       [ -4.716728  ,   3.5809178 ,   8.460974  , ..., -11.874376  ,\n",
       "        -14.738298  , -19.674625  ],\n",
       "       [ 10.3869    ,   0.33998686,  -2.2919376 , ...,  -3.548169  ,\n",
       "         -4.678543  ,   1.2437406 ],\n",
       "       ...,\n",
       "       [  5.6830344 ,  14.351683  ,  46.228428  , ...,  -4.9663544 ,\n",
       "         13.358488  ,   6.82675   ],\n",
       "       [ 25.55392   ,  -7.1336083 ,  10.141382  , ...,   9.633443  ,\n",
       "        -25.31231   ,  -2.9148455 ],\n",
       "       [ 24.146011  ,   6.7507944 ,  18.584644  , ...,  14.319832  ,\n",
       "          3.0579646 ,  29.500877  ]], dtype=float32)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MATRIX MULTIPLICATION PADDED"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "np:  0.0012350082397460938\n",
      "py:  0.0012259483337402344\n",
      "cl:  0.056999921798706055\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "((300, 500), 300, 400, 500)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_times, py_times, cl_times = [], [], []\n",
    "\n",
    "mult = 2\n",
    "M, N, P = 3*(10**mult), 4*(10**mult), 5*(10**mult)\n",
    "M_, N_, P_ = roundUp(M, BLOCK_SIZE), roundUp(N, BLOCK_SIZE), roundUp(P, BLOCK_SIZE)\n",
    "A = np.random.randn(M, N).astype('float32')\n",
    "B = np.random.randn(N, P).astype('float32')\n",
    "ctx = cl.create_some_context()\n",
    "queue = cl.CommandQueue(ctx)\n",
    "mf = cl.mem_flags\n",
    "a_device = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=A)\n",
    "b_device = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=B)\n",
    "\n",
    "prg = cl.Program(ctx, kernel).build()\n",
    "start = time.time()\n",
    "C = A.dot(B)\n",
    "end = time.time()\n",
    "print('np: ', end - start)\n",
    "np_times.append(end - start)\n",
    "start = time.time()\n",
    "C = A @ B\n",
    "end = time.time()\n",
    "print('py: ', end - start)\n",
    "py_times.append(end - start)\n",
    "\n",
    "C_out = np.zeros_like(C)\n",
    "c_device = cl.Buffer(ctx, mf.WRITE_ONLY, C_out.nbytes)\n",
    "C.shape\n",
    "\n",
    "with Timer() as t:\n",
    "    event = prg.GEMM_IMITATE_PADDING(queue, (M_ , P_), (BLOCK_SIZE, BLOCK_SIZE), \n",
    "                     a_device, b_device, c_device,\n",
    "                     np.int32(M), np.int32(N), np.int32(P),\n",
    "                     np.int32(M_), np.int32(N_), np.int32(P_))\n",
    "    event.wait()\n",
    "print('cl: ', t.interval)\n",
    "cl_times.append(t.interval)\n",
    "cl.enqueue_copy(queue, C_out, c_device)\n",
    "\n",
    "C_out.shape, M, N, P"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 10.297104 ,  -1.027521 ,   6.006107 , ..., -19.869083 ,\n",
       "         -1.552961 ,  18.661287 ],\n",
       "       [  1.3731815,  -5.17045  , -16.228472 , ...,  -8.745936 ,\n",
       "         25.594524 ,  15.454743 ],\n",
       "       [ 14.936704 ,  27.659956 ,   0.9331303, ...,  47.217636 ,\n",
       "        -18.784971 ,  -9.16756  ],\n",
       "       ...,\n",
       "       [ -3.0717552,  14.189445 ,  20.689833 , ...,  18.131601 ,\n",
       "         -5.846238 ,  18.857466 ],\n",
       "       [ -1.8767414,  27.955114 ,  18.338612 , ...,  20.37863  ,\n",
       "          4.8617253, -29.914093 ],\n",
       "       [-11.818918 ,  14.612467 ,   9.593132 , ...,   5.0218234,\n",
       "         14.560721 , -11.932028 ]], dtype=float32)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 10.297104  ,  -1.0275171 ,   6.0061083 , ..., -19.86907   ,\n",
       "         -1.5529585 ,  18.661293  ],\n",
       "       [  1.373178  ,  -5.1704493 , -16.228472  , ...,  -8.745934  ,\n",
       "         25.59454   ,  15.454743  ],\n",
       "       [ 14.936706  ,  27.659946  ,   0.93313074, ...,  47.217636  ,\n",
       "        -18.78497   ,  -9.167558  ],\n",
       "       ...,\n",
       "       [ -3.0717509 ,  14.189448  ,  20.689816  , ...,  18.131584  ,\n",
       "         -5.8462415 ,  18.857464  ],\n",
       "       [ -1.8767414 ,  27.955116  ,  18.33861   , ...,  20.37863   ,\n",
       "          4.861725  , -29.914097  ],\n",
       "       [-11.81892   ,  14.612468  ,   9.593132  , ...,   5.0218186 ,\n",
       "         14.560723  , -11.932032  ]], dtype=float32)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MATRIX MULTIPLICATION PADDED REGISTER TILING"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "np:  0.0016679763793945312\n",
      "py:  0.0016739368438720703\n",
      "cl:  0.1386098861694336\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "((300, 500), 300, 400, 500)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np_times, py_times, cl_times = [], [], []\n",
    "\n",
    "ctx = cl.create_some_context()\n",
    "queue = cl.CommandQueue(ctx)\n",
    "mf = cl.mem_flags\n",
    "a_device = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=A)\n",
    "b_device = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=B)\n",
    "\n",
    "prg = cl.Program(ctx, kernel).build()\n",
    "start = time.time()\n",
    "C = A.dot(B)\n",
    "end = time.time()\n",
    "print('np: ', end - start)\n",
    "np_times.append(end - start)\n",
    "start = time.time()\n",
    "C = A @ B\n",
    "end = time.time()\n",
    "print('py: ', end - start)\n",
    "py_times.append(end - start)\n",
    "\n",
    "C_out = np.zeros_like(C)\n",
    "c_device = cl.Buffer(ctx, mf.WRITE_ONLY, C_out.nbytes)\n",
    "C.shape\n",
    "\n",
    "with Timer() as t:\n",
    "    event = prg.GEMM_2DREG_IMITATE_PADDING(queue, (M_ // WPT , P_ // WPT), (BLOCK_SIZE // WPT, BLOCK_SIZE // WPT), \n",
    "                     a_device, b_device, c_device,\n",
    "                     np.int32(M), np.int32(N), np.int32(P),\n",
    "                     np.int32(M_), np.int32(N_), np.int32(P_))\n",
    "    event.wait()\n",
    "print('cl: ', t.interval)\n",
    "cl_times.append(t.interval)\n",
    "cl.enqueue_copy(queue, C_out, c_device)\n",
    "\n",
    "C_out.shape, M, N, P"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 10.297104  ,  -1.0275171 ,   6.0061083 , ..., -19.86907   ,\n",
       "         -1.5529585 ,  18.661293  ],\n",
       "       [  1.373178  ,  -5.1704493 , -16.228472  , ...,  -8.745934  ,\n",
       "         25.59454   ,  15.454743  ],\n",
       "       [ 14.936706  ,  27.659946  ,   0.93313074, ...,  47.217636  ,\n",
       "        -18.78497   ,  -9.167558  ],\n",
       "       ...,\n",
       "       [ -3.0717509 ,  14.189448  ,  20.689816  , ...,  18.131584  ,\n",
       "         -5.8462415 ,  18.857464  ],\n",
       "       [ -1.8767414 ,  27.955116  ,  18.33861   , ...,  20.37863   ,\n",
       "          4.861725  , -29.914097  ],\n",
       "       [-11.81892   ,  14.612468  ,   9.593132  , ...,   5.0218186 ,\n",
       "         14.560723  , -11.932032  ]], dtype=float32)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 10.297104 ,  -1.027521 ,   6.006107 , ..., -19.869083 ,\n",
       "         -1.552961 ,  18.661287 ],\n",
       "       [  1.3731815,  -5.17045  , -16.228472 , ...,  -8.745936 ,\n",
       "         25.594524 ,  15.454743 ],\n",
       "       [ 14.936704 ,  27.659956 ,   0.9331303, ...,  47.217636 ,\n",
       "        -18.784971 ,  -9.16756  ],\n",
       "       ...,\n",
       "       [ -3.0717552,  14.189445 ,  20.689833 , ...,  18.131601 ,\n",
       "         -5.846238 ,  18.857466 ],\n",
       "       [ -1.8767414,  27.955114 ,  18.338612 , ...,  20.37863  ,\n",
       "          4.8617253, -29.914093 ],\n",
       "       [-11.818918 ,  14.612467 ,   9.593132 , ...,   5.0218234,\n",
       "         14.560721 , -11.932028 ]], dtype=float32)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}