mojo-speedtest.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%python\n",
    "def matmul_python(C, A, B):\n",
    "    for m in range(C.rows):\n",
    "        for k in range(A.cols):\n",
    "            for n in range(C.cols):\n",
    "                C[m, n] += A[m, k] * B[k, n]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%python\n",
    "import numpy as np\n",
    "from timeit import timeit\n",
    "\n",
    "class Matrix:\n",
    "    def __init__(self, value, rows, cols):\n",
    "        self.value = value\n",
    "        self.rows = rows\n",
    "        self.cols = cols\n",
    "        \n",
    "    def __getitem__(self, idxs):\n",
    "        return self.value[idxs[0]][idxs[1]]\n",
    "    \n",
    "    def __setitem__(self, idxs, value):\n",
    "        self.value[idxs[0]][idxs[1]] = value\n",
    "\n",
    "def benchmark_matmul_python(M, N, K):\n",
    "    A = Matrix(list(np.random.rand(M, K)), M, K)\n",
    "    B = Matrix(list(np.random.rand(K, N)), K, N)\n",
    "    C = Matrix(list(np.zeros((M, N))), M, N)\n",
    "    secs = timeit(lambda: matmul_python(C, A, B), number=2)/2\n",
    "    gflops = ((2*M*N*K)/secs) / 1e9\n",
    "    print(gflops, \"GFLOP/s\")\n",
    "    return gflops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.0019262170246448786 GFLOP/s\r\n"
     ]
    }
   ],
   "source": [
    "python_gflops = benchmark_matmul_python(128, 128, 128).to_float64()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from benchmark import Benchmark\n",
    "from sys.intrinsics import strided_load\n",
    "from utils.list import VariadicList\n",
    "from math import div_ceil, min\n",
    "from memory import memset_zero\n",
    "from memory.unsafe import DTypePointer\n",
    "from random import rand, random_float64\n",
    "from sys.info import simdwidthof\n",
    "from runtime.llcl import Runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This exactly the same Python implementation, \n",
    "# but is infact Mojo code!\n",
    "def matmul_untyped(C, A, B):\n",
    "    for m in range(C.rows):\n",
    "        for k in range(A.cols):\n",
    "            for n in range(C.cols):\n",
    "                C[m, n] += A[m, k] * B[k, n]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "fn matrix_getitem(self: object, i: object) raises -> object:\n",
    "    return self.value[i]\n",
    "\n",
    "\n",
    "fn matrix_setitem(self: object, i: object, value: object) raises -> object:\n",
    "    self.value[i] = value\n",
    "    return None\n",
    "\n",
    "\n",
    "fn matrix_append(self: object, value: object) raises -> object:\n",
    "    self.value.append(value)\n",
    "    return None\n",
    "\n",
    "\n",
    "fn matrix_init(rows: Int, cols: Int) raises -> object:\n",
    "    let value = object([])\n",
    "    return object(\n",
    "        Attr(\"value\", value), Attr(\"__getitem__\", matrix_getitem), Attr(\"__setitem__\", matrix_setitem), \n",
    "        Attr(\"rows\", rows), Attr(\"cols\", cols), Attr(\"append\", matrix_append),\n",
    "    )\n",
    "\n",
    "def benchmark_matmul_untyped(M: Int, N: Int, K: Int, python_gflops: Float64):\n",
    "    C = matrix_init(M, N)\n",
    "    A = matrix_init(M, K)\n",
    "    B = matrix_init(K, N)\n",
    "    for i in range(M):\n",
    "        c_row = object([])\n",
    "        b_row = object([])\n",
    "        a_row = object([])\n",
    "        for j in range(N):\n",
    "            c_row.append(0.0)\n",
    "            b_row.append(random_float64(-5, 5))\n",
    "            a_row.append(random_float64(-5, 5))\n",
    "        C.append(c_row)\n",
    "        B.append(b_row)\n",
    "        A.append(a_row)\n",
    "\n",
    "    @parameter\n",
    "    fn test_fn():\n",
    "        try:\n",
    "            _ = matmul_untyped(C, A, B)\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "    let secs = Float64(Benchmark().run[test_fn]()) / 1_000_000_000\n",
    "    _ = (A, B, C)\n",
    "    let gflops = ((2*M*N*K)/secs) / 1e9\n",
    "    let speedup : Float64 = gflops / python_gflops\n",
    "    print(gflops, \"GFLOP/s, a\", speedup.value, \"x speedup over Python\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.015576809199651842 GFLOP/s, a 8.0867363336297036 x speedup over Python\r\n"
     ]
    }
   ],
   "source": [
    "benchmark_matmul_untyped(128, 128, 128, python_gflops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct Matrix:\n",
    "    var data: DTypePointer[DType.float32]\n",
    "    var rows: Int\n",
    "    var cols: Int\n",
    "\n",
    "    fn __init__(inout self, rows: Int, cols: Int):\n",
    "        self.data = DTypePointer[DType.float32].alloc(rows * cols)\n",
    "        rand(self.data, rows*cols)\n",
    "        self.rows = rows\n",
    "        self.cols = cols\n",
    "\n",
    "    fn __del__(owned self):\n",
    "        self.data.free()\n",
    "\n",
    "    fn zero(inout self):\n",
    "        memset_zero(self.data, self.rows * self.cols)\n",
    "\n",
    "    @always_inline\n",
    "    fn __getitem__(self, y: Int, x: Int) -> Float32:\n",
    "        return self.load[1](y, x)\n",
    "\n",
    "    @always_inline\n",
    "    fn load[nelts:Int](self, y: Int, x: Int) -> SIMD[DType.float32, nelts]:\n",
    "        return self.data.simd_load[nelts](y * self.cols + x)\n",
    "\n",
    "    @always_inline\n",
    "    fn __setitem__(self, y: Int, x: Int, val: Float32):\n",
    "        return self.store[1](y, x, val)\n",
    "\n",
    "    @always_inline\n",
    "    fn store[nelts:Int](self, y: Int, x: Int, val: SIMD[DType.float32, nelts]):\n",
    "        self.data.simd_store[nelts](y * self.cols + x, val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note that C, A, and B have types.\n",
    "fn matmul_naive(C: Matrix, A: Matrix, B: Matrix):\n",
    "    for m in range(C.rows):\n",
    "        for k in range(A.cols):\n",
    "            for n in range(C.cols):\n",
    "                C[m, n] += A[m, k] * B[k, n]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "@always_inline\n",
    "fn benchmark[\n",
    "    func: fn (Matrix, Matrix, Matrix) -> None\n",
    "](M: Int, N: Int, K: Int, base_gflops: Float64):\n",
    "    var C = Matrix(M, N)\n",
    "    C.zero()\n",
    "    var A = Matrix(M, K)\n",
    "    var B = Matrix(K, N)\n",
    "\n",
    "    @always_inline\n",
    "    @parameter\n",
    "    fn test_fn():\n",
    "        _ = func(C, A, B)\n",
    "\n",
    "    let secs = Float64(Benchmark().run[test_fn]()) / 1_000_000_000\n",
    "    # Prevent the matrices from being freed before the benchmark run\n",
    "    _ = (A, B, C)\n",
    "    let gflops = ((2 * M * N * K) / secs) / 1e9\n",
    "    let speedup: Float64 = gflops / base_gflops\n",
    "    # print(gflops, \"GFLOP/s\", speedup, \" speedup\")\n",
    "    print(gflops, \"GFLOP/s, a\", speedup.value, \"x speedup over Python\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.9400845456250755 GFLOP/s, a 3083.8085582387589 x speedup over Python\r\n"
     ]
    }
   ],
   "source": [
    "benchmark[matmul_naive](512, 512, 512, python_gflops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mojo has SIMD vector types, we can vectorize the Matmul code as follows.\n",
    "alias nelts = simdwidthof[DType.float32]() # The SIMD vector width.\n",
    "fn matmul_vectorized_0(C: Matrix, A: Matrix, B: Matrix):\n",
    "    for m in range(C.rows):\n",
    "        for k in range(A.cols):\n",
    "            for nv in range(0, C.cols, nelts):\n",
    "                C.store[nelts](m,nv, C.load[nelts](m,nv) + A[m,k] * B.load[nelts](k,nv))\n",
    "        \n",
    "            # Handle remaining elements with scalars.\n",
    "            for n in range(nelts*(C.cols//nelts), C.cols):\n",
    "                C[m,n] += A[m,k] * B[k,n]\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "25.630211392670954 GFLOP/s, a 13305.983212040292 x speedup over Python\r\n"
     ]
    }
   ],
   "source": [
    "benchmark[matmul_vectorized_0](512, 512, 512, python_gflops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simplify the code by using the builtin vectorize function\n",
    "from algorithm import vectorize\n",
    "fn matmul_vectorized_1(C: Matrix, A: Matrix, B: Matrix):\n",
    "    for m in range(C.rows):\n",
    "        for k in range(A.cols):\n",
    "            @parameter\n",
    "            fn dot[nelts : Int](n : Int):\n",
    "                C.store[nelts](m,n, C.load[nelts](m,n) + A[m,k] * B.load[nelts](k,n))\n",
    "            vectorize[nelts, dot](C.cols)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "25.42251593787838 GFLOP/s, a 13198.15763883892 x speedup over Python\r\n"
     ]
    }
   ],
   "source": [
    "benchmark[matmul_vectorized_1](512, 512, 512, python_gflops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "@always_inline\n",
    "fn benchmark_parallel[\n",
    "    func: fn (Matrix, Matrix, Matrix, Runtime) -> None\n",
    "](M: Int, N: Int, K: Int, base_gflops: Float64):\n",
    "    var C = Matrix(M, N)\n",
    "    C.zero()\n",
    "    var A = Matrix(M, K)\n",
    "    var B = Matrix(K, N)\n",
    "\n",
    "    with Runtime() as rt:\n",
    "        @always_inline\n",
    "        @parameter\n",
    "        fn test_fn():\n",
    "            _ = func(C, A, B, rt)\n",
    "\n",
    "        let secs = Float64(Benchmark().run[test_fn]()) / 1_000_000_000\n",
    "        # Prevent the matrices from being freed before the benchmark run\n",
    "        _ = (A, B, C)\n",
    "        let gflops = ((2 * M * N * K) / secs) / 1e9\n",
    "        let speedup: Float64 = gflops / base_gflops\n",
    "        # print(gflops, \"GFLOP/s\", speedup, \" speedup\")\n",
    "        print(gflops, \"GFLOP/s, a\", speedup.value, \"x speedup over Python\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parallelize the code by using the builtin parallelize function\n",
    "from algorithm import parallelize\n",
    "fn matmul_parallelized(C: Matrix, A: Matrix, B: Matrix, rt: Runtime):\n",
    "    @parameter\n",
    "    fn calc_row(m: Int):\n",
    "        for k in range(A.cols):\n",
    "            @parameter\n",
    "            fn dot[nelts : Int](n : Int):\n",
    "                C.store[nelts](m,n, C.load[nelts](m,n) + A[m,k] * B.load[nelts](k,n))\n",
    "            vectorize[nelts, dot](C.cols)\n",
    "        \n",
    "    parallelize[calc_row](rt, C.rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "43.067415829484553 GFLOP/s, a 22358.548013261669 x speedup over Python\r\n"
     ]
    }
   ],
   "source": [
    "benchmark_parallel[matmul_parallelized](512, 512, 512, python_gflops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from algorithm import Static2DTileUnitFunc as Tile2DFunc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Perform 2D tiling on the iteration space defined by end_x and end_y.\n",
    "fn tile[tiled_fn: Tile2DFunc, tile_x: Int, tile_y: Int](end_x: Int, end_y: Int):\n",
    "    # Note: this assumes that ends are multiples of the tiles.\n",
    "    for y in range(0, end_y, tile_y):\n",
    "        for x in range(0, end_x, tile_x):\n",
    "            tiled_fn[tile_x, tile_y](x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the above tile function to perform tiled matmul.\n",
    "fn matmul_tiled_parallelized(C: Matrix, A: Matrix, B: Matrix, rt: Runtime):\n",
    "    @parameter\n",
    "    fn calc_row(m: Int):\n",
    "        @parameter\n",
    "        fn calc_tile[tile_x: Int, tile_y: Int](x: Int, y: Int):\n",
    "            for k in range(y, y + tile_y):\n",
    "                @parameter\n",
    "                fn dot[nelts : Int,](n : Int):\n",
    "                    C.store[nelts](m,n + x, C.load[nelts](m,n+x) + A[m,k] * B.load[nelts](k,n+x))\n",
    "                vectorize[nelts, dot](tile_x)\n",
    "\n",
    "        # We hardcode the tile factor to be 4.\n",
    "        alias tile_size = 4\n",
    "        tile[calc_tile, nelts * tile_size, tile_size](A.cols, C.cols)\n",
    "\n",
    "    parallelize[calc_row](rt, C.rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "38.069228341945291 GFLOP/s, a 19763.727479754682 x speedup over Python\r\n"
     ]
    }
   ],
   "source": [
    "benchmark_parallel[matmul_tiled_parallelized](512, 512, 512, python_gflops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Unroll the vectorized loop by a constant factor.\n",
    "from algorithm import vectorize_unroll\n",
    "fn matmul_tiled_unrolled_parallelized(C: Matrix, A: Matrix, B: Matrix, rt: Runtime):\n",
    "    @parameter\n",
    "    fn calc_row(m: Int):\n",
    "        @parameter\n",
    "        fn calc_tile[tile_x: Int, tile_y: Int](x: Int, y: Int):\n",
    "            for k in range(y, y + tile_y):\n",
    "                @parameter\n",
    "                fn dot[nelts : Int,](n : Int):\n",
    "                    C.store[nelts](m,n+x, C.load[nelts](m,n+x) + A[m,k] * B.load[nelts](k,n+x))\n",
    "\n",
    "                # Vectorize by nelts and unroll by tile_x/nelts\n",
    "                # Here unroll factor is 4\n",
    "                vectorize_unroll[nelts, tile_x//nelts, dot](tile_x)\n",
    "\n",
    "        alias tile_size = 4\n",
    "        tile[calc_tile, nelts*tile_size, tile_size](A.cols, C.cols)\n",
    "      \n",
    "    parallelize[calc_row](rt, C.rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "40.100317623076329 GFLOP/s, a 20818.172152989515 x speedup over Python\r\n"
     ]
    }
   ],
   "source": [
    "benchmark_parallel[matmul_tiled_unrolled_parallelized](512, 512, 512, python_gflops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autotune import autotune, search\n",
    "from time import now\n",
    "from memory.unsafe import Pointer\n",
    "\n",
    "alias matmul_fn_sig_type = fn(Matrix, Matrix, Matrix, Runtime) -> None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Autotune the tile size used in the matmul.\n",
    "@adaptive\n",
    "fn matmul_autotune_impl(C: Matrix, A: Matrix, B: Matrix, rt: Runtime):\n",
    "    @parameter\n",
    "    fn calc_row(m: Int):\n",
    "        @parameter\n",
    "        fn calc_tile[tile_x: Int, tile_y: Int](x: Int, y: Int):\n",
    "            for k in range(y, y + tile_y):\n",
    "                @parameter\n",
    "                fn dot[nelts : Int,](n : Int):\n",
    "                    C.store[nelts](m,n+x, C.load[nelts](m,n+x) + A[m,k] * B.load[nelts](k,n+x))\n",
    "                vectorize_unroll[nelts, tile_x // nelts, dot](tile_x)\n",
    "\n",
    "        # Instead of hardcoding to tile_size = 4, search for the fastest \n",
    "        # tile size by evaluting this function as tile size varies.\n",
    "        alias tile_size = autotune(1, 2, 4, 8, 16, 32)\n",
    "        tile[calc_tile, nelts * tile_size, tile_size](A.cols, C.cols)\n",
    "      \n",
    "    parallelize[calc_row](rt, C.rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "fn matmul_evaluator(funcs: Pointer[matmul_fn_sig_type], size: Int) -> Int:\n",
    "    print(\"matmul_evaluator, number of candidates: \", size)\n",
    "\n",
    "    let eval_begin: Int = now()\n",
    "\n",
    "    # This size is picked at random, in real code we could use a real size\n",
    "    # distribution here.\n",
    "    let M = 512\n",
    "    let N = 512\n",
    "    let K = 512\n",
    "    print(\"Optimizing for size:\", M, \"x\", N, \"x\", K)\n",
    "\n",
    "    var best_idx: Int = -1\n",
    "    var best_time: Int = -1\n",
    "\n",
    "    alias eval_iterations = 10\n",
    "    alias eval_samples = 10\n",
    "\n",
    "    var C = Matrix(M, N)\n",
    "    var A = Matrix(M, K)\n",
    "    var B = Matrix(K, N)\n",
    "    let Cptr = Pointer[Matrix].address_of(C).address\n",
    "    let Aptr = Pointer[Matrix].address_of(A).address\n",
    "    let Bptr = Pointer[Matrix].address_of(B).address\n",
    "    with Runtime() as rt:\n",
    "        # Find the function that's the fastest on the size we're optimizing for\n",
    "        for f_idx in range(size):\n",
    "            let func = funcs.load(f_idx)\n",
    "\n",
    "            @always_inline\n",
    "            @parameter\n",
    "            fn wrapper():\n",
    "                func(C, A, B, rt)\n",
    "            let cur_time = Benchmark(1, 100_000, 500_000_000, 1000_000_000).run[wrapper]()\n",
    "\n",
    "            if best_idx < 0:\n",
    "                best_idx = f_idx\n",
    "                best_time = cur_time\n",
    "            if best_time > cur_time:\n",
    "                best_idx = f_idx\n",
    "                best_time = cur_time\n",
    "\n",
    "        let eval_end: Int = now()\n",
    "        # Prevent matrices from being destroyed before we finished benchmarking them.\n",
    "        _ = A.data\n",
    "        _ = B.data\n",
    "        _ = C.data\n",
    "        print(\"Time spent in matmul_evaluator, ms:\", (eval_end - eval_begin) // 1000000)\n",
    "        print(\"Best candidate idx:\", best_idx)\n",
    "        return best_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "fn matmul_autotune(C: Matrix, A: Matrix, B: Matrix, rt: Runtime):\n",
    "    alias best_impl: matmul_fn_sig_type\n",
    "    search[\n",
    "        matmul_fn_sig_type,\n",
    "        VariadicList(matmul_autotune_impl.__adaptive_set),\n",
    "        matmul_evaluator -> best_impl\n",
    "    ]()\n",
    "    # Run the best candidate\n",
    "    return best_impl(C, A, B, rt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "matmul_evaluator, number of candidates:  6\n",
      "Optimizing for size: 512 x 512 x 512\n",
      "Time spent in matmul_evaluator, ms: 7896\n",
      "Best candidate idx: 3\n",
      "40.366573914005897 GFLOP/s, a 20956.399719002569 x speedup over Python\n"
     ]
    }
   ],
   "source": [
    "benchmark_parallel[matmul_autotune](512, 512, 512, python_gflops)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Mojo",
   "language": "mojo",
   "name": "mojo-jupyter-kernel"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "mojo"
   },
   "file_extension": ".mojo",
   "mimetype": "text/x-mojo",
   "name": "mojo"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}