llama2_gemm_tiling_analysis_ffn_up.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "from gemm_tiling import gemm_tiling_input_stationary, gemm_tiling_weight_stationary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FFN Up analysis, B=1\n",
      "FFN Up, M=4096, K=4096, N=11008, B=1, tile_m=128, tile_n=32, stationary: input, utilization=89.10%\n",
      "Warning: using defautl buffer strategy, nedd total memory 3.028809 MB, > SRAM 3.0 MB\n",
      "Warning: change input buffer strategy to input_buffer_num =  1\n",
      "FFN Up, M=4096, K=4096, N=11008, B=1, tile_m=128, tile_n=43, stationary: input, utilization=91.17%\n",
      "Warning: using defautl buffer strategy, nedd total memory 4.781250 MB, > SRAM 3.0 MB\n",
      "Warning: change input buffer strategy to input_buffer_num =  1\n",
      "FFN Up, M=4096, K=4096, N=11008, B=1, tile_m=256, tile_n=32, stationary: input, utilization=87.26%\n",
      "Warning: using defautl buffer strategy, nedd total memory 4.512695 MB, > SRAM 3.0 MB\n",
      "Warning: change input buffer strategy to input_buffer_num =  1\n",
      "FFN Up, M=4096, K=4096, N=11008, B=1, tile_m=256, tile_n=21, stationary: input, utilization=88.88%\n",
      "FFN Up, M=4096, K=4096, N=11008, B=1, tile_m=32, tile_n=128, stationary: weight, utilization=74.68%\n",
      "Warning: using defautl buffer strategy, nedd total memory 4.781250 MB, > SRAM 3.0 MB\n",
      "Warning: change input buffer strategy to input_buffer_num =  1\n",
      "FFN Up, M=4096, K=4096, N=11008, B=1, tile_m=32, tile_n=256, stationary: weight, utilization=73.31%\n"
     ]
    }
   ],
   "source": [
    "# =============== FNN Up  ===================\n",
    "# FFN Gate 与 FFN Up是完全同样的size的计算，这里不列出了\n",
    "M, K, N = 4096, 4096, 11008\n",
    "B = 1\n",
    "\n",
    "print(f\"FFN Up analysis, B={B}\")\n",
    "\n",
    "tile_m = 128\n",
    "tile_n = 32\n",
    "utilization = gemm_tiling_input_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: input, utilization={utilization:.2f}%\")\n",
    "\n",
    "\n",
    "tile_m = 128\n",
    "tile_n = 43\n",
    "utilization = gemm_tiling_input_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: input, utilization={utilization:.2f}%\")\n",
    "\n",
    "\n",
    "tile_m = 256\n",
    "tile_n = 32\n",
    "utilization = gemm_tiling_input_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: input, utilization={utilization:.2f}%\")\n",
    "\n",
    "tile_m = 256\n",
    "tile_n = 21\n",
    "utilization = gemm_tiling_input_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: input, utilization={utilization:.2f}%\")\n",
    "\n",
    "tile_m = 32\n",
    "tile_n = 128\n",
    "utilization = gemm_tiling_weight_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: weight, utilization={utilization:.2f}%\")\n",
    "\n",
    "tile_m = 32\n",
    "tile_n = 256\n",
    "utilization = gemm_tiling_weight_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: weight, utilization={utilization:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "因为 FFN up是升维，N的维度很高11008，所以input stationary的性能会比weight stationary好很多\n",
    "\n",
    "下面分析batch size的影响"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FFN Up analysis, B=16\n",
      "FFN Up, M=4096, K=4096, N=11008, B=16, tile_m=128, tile_n=32, stationary: input, utilization=91.96%\n",
      "Warning: using defautl buffer strategy, nedd total memory 3.028809 MB, > SRAM 3.0 MB\n",
      "Warning: change input buffer strategy to input_buffer_num =  1\n",
      "FFN Up, M=4096, K=4096, N=11008, B=16, tile_m=128, tile_n=43, stationary: input, utilization=95.17%\n",
      "Warning: using defautl buffer strategy, nedd total memory 4.781250 MB, > SRAM 3.0 MB\n",
      "Warning: change input buffer strategy to input_buffer_num =  1\n",
      "FFN Up, M=4096, K=4096, N=11008, B=16, tile_m=256, tile_n=32, stationary: input, utilization=91.78%\n",
      "Warning: using defautl buffer strategy, nedd total memory 4.512695 MB, > SRAM 3.0 MB\n",
      "Warning: change input buffer strategy to input_buffer_num =  1\n",
      "FFN Up, M=4096, K=4096, N=11008, B=16, tile_m=256, tile_n=21, stationary: input, utilization=91.91%\n",
      "FFN Up, M=4096, K=4096, N=11008, B=16, tile_m=32, tile_n=128, stationary: weight, utilization=88.48%\n",
      "Warning: using defautl buffer strategy, nedd total memory 4.781250 MB, > SRAM 3.0 MB\n",
      "Warning: change input buffer strategy to input_buffer_num =  1\n",
      "FFN Up, M=4096, K=4096, N=11008, B=16, tile_m=32, tile_n=256, stationary: weight, utilization=88.36%\n",
      "FFN Up, M=4096, K=4096, N=11008, B=16, tile_m=8, tile_n=128, stationary: weight, utilization=88.11%\n"
     ]
    }
   ],
   "source": [
    "M, K, N = 4096, 4096, 11008\n",
    "B = 16\n",
    "print(f\"FFN Up analysis, B={B}\")\n",
    "\n",
    "\n",
    "tile_m = 128\n",
    "tile_n = 32\n",
    "utilization = gemm_tiling_input_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: input, utilization={utilization:.2f}%\")\n",
    "\n",
    "\n",
    "tile_m = 128\n",
    "tile_n = 43\n",
    "utilization = gemm_tiling_input_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: input, utilization={utilization:.2f}%\")\n",
    "\n",
    "\n",
    "tile_m = 256\n",
    "tile_n = 32\n",
    "utilization = gemm_tiling_input_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: input, utilization={utilization:.2f}%\")\n",
    "\n",
    "tile_m = 256\n",
    "tile_n = 21\n",
    "utilization = gemm_tiling_input_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: input, utilization={utilization:.2f}%\")\n",
    "\n",
    "tile_m = 32\n",
    "tile_n = 128\n",
    "utilization = gemm_tiling_weight_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: weight, utilization={utilization:.2f}%\")\n",
    "\n",
    "tile_m = 32\n",
    "tile_n = 256\n",
    "utilization = gemm_tiling_weight_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: weight, utilization={utilization:.2f}%\")\n",
    "\n",
    "tile_m = 8\n",
    "tile_n = 128\n",
    "utilization = gemm_tiling_weight_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: weight, utilization={utilization:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FFN Up, M=4096, K=4096, N=11008, B=1, tile_m=4, tile_n=86, stationary: weight, utilization=78.16%\n",
      "FFN Up, M=4096, K=4096, N=11008, B=16, tile_m=4, tile_n=86, stationary: weight, utilization=98.28%\n"
     ]
    }
   ],
   "source": [
    "M, K, N = 4096, 4096, 11008\n",
    "B = 1\n",
    "tile_m, tile_n = 4, 86\n",
    "utilization = gemm_tiling_weight_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: weight, utilization={utilization:.2f}%\")\n",
    "B = 16\n",
    "tile_m, tile_n = 4, 86\n",
    "utilization = gemm_tiling_weight_stationary(B, M, K, N, tile_m, tile_n, print_details=False)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: weight, utilization={utilization:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "详细分析"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=========memory anylasis ===================\n",
      "+-------------------+--------------------------+\n",
      "|        var        |         mem (MB)         |\n",
      "+-------------------+--------------------------+\n",
      "|    input_size     | 0.671875 * 2 =  1.343750 |\n",
      "|    weight_size    | 0.031250 * 3 =  0.093750 |\n",
      "|    output_size    | 0.000656 * 2 =  0.001312 |\n",
      "|    total_size     |         1.438812         |\n",
      "| input_load_iters  |            8             |\n",
      "| weight_load_iters |           1024           |\n",
      "+-------------------+--------------------------+\n",
      "=========unit time anylasis ===================\n",
      "+-----------------------+------------+\n",
      "|         unit          | time (us)  |\n",
      "+-----------------------+------------+\n",
      "|    input_load_time    | 105.080469 |\n",
      "|   weight_load_time    |  4.982813  |\n",
      "|    weight_noc_time    |  0.248419  |\n",
      "| compute_time_one_tile |  0.352256  |\n",
      "|   output_save_time    |  0.202520  |\n",
      "+-----------------------+------------+\n",
      "=========internal time anylasis ===================\n",
      "+--------------------------+------------+\n",
      "|           item           | time (us)  |\n",
      "+--------------------------+------------+\n",
      "|  time_one_noc_pipe_flow  |  5.636096  |\n",
      "|     time_one_iter_w      |  5.636096  |\n",
      "|     time_one_iter_in     | 105.080469 |\n",
      "| compute_time_in_one_iter |  5.636096  |\n",
      "+--------------------------+------------+\n",
      "=========final time anylasis ===================\n",
      "+-----------------------+--------------+------------+\n",
      "|         item          |  time (us)   | percentage |\n",
      "+-----------------------+--------------+------------+\n",
      "|   inital_load_time    |  110.063281  |   0.23%    |\n",
      "| iter_over_weight_time | 46125.809664 |   98.19%   |\n",
      "| iter_over_input_time  |  735.563281  |   1.57%    |\n",
      "|    last_iter_time     |   5.636096   |   0.01%    |\n",
      "|      total_time       | 46977.274842 |    100%    |\n",
      "|  total_compute_time   | 46170.898432 |   98.28%   |\n",
      "|      utilization      |    98.28%    |            |\n",
      "+-----------------------+--------------+------------+\n",
      "FFN Up, M=4096, K=4096, N=11008, B=16, tile_m=4, tile_n=86, stationary: weight, utilization=98.28%\n"
     ]
    }
   ],
   "source": [
    "M, K, N = 4096, 4096, 11008\n",
    "B = 16\n",
    "tile_m, tile_n = 4, 86\n",
    "utilization = gemm_tiling_weight_stationary(B, M, K, N, tile_m, tile_n, print_details=True)\n",
    "print(f\"FFN Up, M={M}, K={K}, N={N}, B={B}, tile_m={tile_m}, tile_n={tile_n}, stationary: weight, utilization={utilization:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "# constants\n",
    "GB = 1024*1024*1024\n",
    "MB = 1024*1024\n",
    "KB = 1024\n",
    "ns = 1e-9\n",
    "us = 1e-6\n",
    "ms = 1e-3\n",
    "# 1 TFLOPS = 1e12 FLOPS 还是 2**40 FLOPS （1.0995e12)? 这个影响还是有一些的\n",
    "TFLOPS = 1e12\n",
    "# hardware configuration\n",
    "data_type = 2  # 2 bytes for FP16, 4 bytes for FP32\n",
    "Tile_num = 4 * 4  # 4x4 tiles\n",
    "Tile_SRAM = 3 * MB  # 3MB\n",
    "Tile_compute = 128/16 * TFLOPS # 8 TFLOPS\n",
    "DDR_BW = 100 * GB  # 100GB/s\n",
    "NOC_BW = 128 * GB  # 128GB/s\n",
    "NOC_latency_hop = 10 * ns  # 10ns for 1 hop\n",
    "DDR_latency = 0 * ns  # 100ns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tile_m need to be > 58.20766091346741, to hide the NOC time\n",
      "tile_m need to be > 74.50580596923828, to hide the DDR time\n"
     ]
    }
   ],
   "source": [
    "print(f\"tile_m need to be > {Tile_compute/NOC_BW}, to hide the NOC time\")\n",
    "\n",
    "print(f\"tile_m need to be > {Tile_compute/DDR_BW}, to hide the DDR time\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "688.0"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "11008/16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "43.0"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "688.0/16"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}