[ { "path": ".gitignore", "content": "__pycache__/\nSlurmOutput/\ntemp/\n.vscode/\n*.ncu-rep\n" }, { "path": ".gitmodules", "content": "[submodule \"cost_model/supply_chain\"]\n\tpath = cost_model/supply_chain\n\turl = https://github.com/PrincetonUniversity/ttm-cas.git\n" }, { "path": "Dockerfile", "content": "# Start with a base image that includes Miniconda to manage our environment\nFROM continuumio/miniconda3\n\n# Set the working directory in the container to /app\nWORKDIR /app\n\n# Create the conda environment\nCOPY environment.yml /app/environment.yml\nRUN conda env create -f /app/environment.yml\n\n# Initialize conda in bash shell\nRUN echo \"source activate llmcompass_ae\" > ~/.bashrc\nENV PATH /opt/conda/envs/llmcompass_ae/bin:$PATH\n\n# Clone your GitHub repository\nRUN git clone https://github.com/HenryChang213/LLMCompass_ISCA_AE.git /app/LLMCompass_ISCA_AE\nRUN cd /app/LLMCompass_ISCA_AE && git submodule init && git submodule update --recursive\n\n# Expose the port your app runs on\nEXPOSE 8000\n\n\n" }, { "path": "LICENSE", "content": "BSD 3-Clause License\n\nCopyright (c) 2024, Princeton University\n\nRedistribution and use in source and binary forms, with or without\nmodification, are permitted provided that the following conditions are met:\n\n1. Redistributions of source code must retain the above copyright notice, this\n list of conditions and the following disclaimer.\n\n2. Redistributions in binary form must reproduce the above copyright notice,\n this list of conditions and the following disclaimer in the documentation\n and/or other materials provided with the distribution.\n\n3. Neither the name of the copyright holder nor the names of its\n contributors may be used to endorse or promote products derived from\n this software without specific prior written permission.\n\nTHIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\nAND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\nIMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\nDISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\nFOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\nDAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\nSERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\nCAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\nOR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\nOF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n" }, { "path": "README.md", "content": "[![DOI](https://zenodo.org/badge/779008229.svg)](https://zenodo.org/doi/10.5281/zenodo.10892431)\n\n# LLMCompass\n\nThis repository provides the implementation of **LLMCompass** from the following papers:\n\n[**LLMCompass: Enabling Efficient Hardware Design for Large Language Model Inference**](https://parallel.princeton.edu/papers/isca24_llmcompass.pdf)\n\n*Hengrui Zhang, August Ning, Rohan Baskar Prabhakar, David Wentzlaff*\n\n\nIn the Proceedings of the 51st Annual International Symposium on Computer Architecture:\n\n```\n@inproceedings{LLMCompass,\nauthor = {Zhang, Hengrui and Ning, August and Prabhakar, Rohan Baskar and Wentzlaff, David},\ntitle = {LLMCompass: Enabling Efficient Hardware Design for Large Language Model Inference},\nyear = {2024},\nbooktitle = {Proceedings of the 51st Annual International Symposium on Computer Architecture},\n}\n```\n\n\n## Set up the environment\n\n```\n$ conda create -n llmcompass_ae python=3.9\n$ conda activate llmcompass_ae\n$ pip3 install scalesim\n$ conda install pytorch==2.0.0 -c pytorch\n$ pip3 install matplotlib\n$ pip3 install seaborn\n$ pip3 install scipy\n```\n\n## Installation\n\n### If using Github\n```\n$ git clone -b ISCA_AE https://github.com/PrincetonUniversity/LLMCompass\n$ cd LLMCompass\n$ git submodule init\n$ git submodule update --recursive\n```\n\n### If using Zenodo\nUnzip the file and download from https://github.com/PrincetonUniversity/ttm-cas.git as `cost_model\\supply_chain`\n\n\n### If using Docker\nA Dockerfile has been provided (`./Dockerfile`), including all the software dependencies and the LLMCompass source code.\n\nA docker image has been provided [here](https://github.com/HenryChang213/LLMCompass_ISCA_AE_docker).\n\n## AE Experiment workflow\n```\n# Figure 5 (around 100 min) \n$ cd ae/figure5\n$ bash run_figure5.sh \n\n# Figure 6 (around 1 min)\n$ cd ae/figure6\n$ bash run_figure6.sh\n\n# Figure 7 (around 20 min)\n$ cd ae/figure7\n$ bash run_figure7.sh\n\n# Figure 8 (around 40 min)\n$ cd ae/figure8\n$ bash run_figure8.sh\n\n# Figure 9 (around 30 min)\n$ cd ae/figure9\n$ bash run_figure9.sh\n\n# Figure 10 (around 45 min)\n$ cd ae/figure10\n$ bash run_figure10.sh\n\n# Figure 11 (around 5 min) \n$ cd ae/figure11\n$ bash run_figure11.sh\n\n# Figure 12 (around 4 hours) \n$ cd ae/figure12\n$ bash run_figure12.sh\n```\n\n## AE Expected result\n\nAfter running each script above, the corresponding figures\nwill be generated under the corresponding directory as suggested by its name.\n\nFor comparison, a copy of the expected results can be found in `ae\\expected_results`\n\n\n## User Guide\n\nA guide on \"How to Run a LLMCompass Simulation\" is shown [here](./docs/run.md)." }, { "path": "__init__.py", "content": "" }, { "path": "ae/.gitignore", "content": "*.pdf\n*.csv\n!**/real_hardware/**/*.csv\n!expected_results/*" }, { "path": "ae/__init__.py", "content": "" }, { "path": "ae/figure10/__init__.py", "content": "" }, { "path": "ae/figure10/plot_latency.py", "content": "import pandas as pd\nimport numpy as np\nimport seaborn as sns\nimport matplotlib.pyplot as plt\n\n\nour_decoding = pd.read_csv(\n \"our_decoding.csv\", header=None, names=[\"bs\", \"s\", \"latency\"]\n).sort_values(by=\"s\")\nour_prefill = pd.read_csv(\"our_prefill.csv\", header=None, names=[\"bs\", \"s\", \"latency\"])\nA100_decoding = pd.read_csv(\n \"A100_decoding.csv\", header=None, names=[\"bs\", \"s\", \"latency\"]\n).sort_values(by=\"s\")\nA100_prefill = pd.read_csv(\n \"A100_prefill.csv\", header=None, names=[\"bs\", \"s\", \"latency\"]\n)\n\n\ndef get_total_decoding_latency(df: pd.DataFrame, start, end):\n df_filtered = df[(df[\"s\"] >= start) & (df[\"s\"] <= end)]\n total_latency = 0\n\n # Calculate the mean of the values for each length interval and add to the sum\n for i in range(len(df_filtered) - 1):\n # Calculate the mean of current and next value\n mean = (df_filtered.iloc[i][\"latency\"] + df_filtered.iloc[i + 1][\"latency\"]) / 2\n # Calculate the difference in length\n length_interval = df_filtered.iloc[i + 1][\"s\"] - df_filtered.iloc[i][\"s\"]\n # Multiply the mean value by the length interval and add to the sum\n total_latency += mean * length_interval\n\n # print(total_latency)\n return total_latency\n\n\nnorm_perf = []\nfor input_length in [256, 512, 1024, 2048]:\n temp_list = []\n our_prefill_latency = our_prefill[our_prefill[\"s\"] == input_length][\n \"latency\"\n ].values[0]\n A100_prefill_latency = A100_prefill[A100_prefill[\"s\"] == input_length][\n \"latency\"\n ].values[0]\n for output_length in [256, 512, 768, 1024, 1280, 1536, 1792, 2048]:\n our_total_latency = our_prefill_latency + get_total_decoding_latency(\n our_decoding, input_length, input_length + output_length\n )\n A100_total_latency = A100_prefill_latency + get_total_decoding_latency(\n A100_decoding, input_length, input_length + output_length\n )\n temp_list.append(A100_total_latency / our_total_latency)\n norm_perf.append(temp_list)\n\ncmap = sns.color_palette(\"viridis\", as_cmap=True)\ndata = np.array(norm_perf)\nimport statistics\n\nprint(statistics.geometric_mean(data.flatten()))\nfig, ax = plt.subplots()\ncax = ax.imshow(data, interpolation=\"nearest\", cmap=cmap, vmin=0.8, vmax=1)\n# cax = sns.heatmap(data, cmap=\"viridis\")\n\n# Add a colorbar\nfig.colorbar(cax, shrink=0.5)\n\n\n# Function to convert RGB to grayscale intensity\ndef get_intensity(color):\n return color[0] * 0.299 + color[1] * 0.587 + color[2] * 0.114\n\n\n# Set a threshold for deciding text color\nintensity_threshold = 0.5\nfor i in range(data.shape[0]):\n for j in range(data.shape[1]):\n # Get the color from the colormap\n cell_color = cax.cmap(cax.norm(data[i, j]))\n # Calculate intensity of the cell color\n intensity = get_intensity(cell_color)\n # Choose text color based on intensity\n text_color = \"white\" if intensity < intensity_threshold else \"black\"\n text = ax.text(\n j, i, round(data[i, j], 2), ha=\"center\", va=\"center\", color=text_color\n )\n\n# Set the x-axis and y-axis values\nx_axis_labels = [256, 512, 768, 1024, 1280, 1536, 1792, 2048]\ny_axis_labels = [256, 512, 1024, 2048]\n\n# Set ticks positions\nax.set_xticks(np.arange(len(x_axis_labels)))\nax.set_yticks(np.arange(len(y_axis_labels)))\n\n# Set ticks labels\nax.set_xticklabels(x_axis_labels)\nax.set_yticklabels(y_axis_labels)\n\n# Set labels for axes\nax.set_xlabel(\"Output Length\")\nax.set_ylabel(\"Input Length\")\nax.invert_yaxis()\n# # Rotate the tick labels for the x-axis if needed\n# plt.setp(ax.get_xticklabels(), rotation=45, ha=\"right\", rotation_mode=\"anchor\")\n\n# Show the plot\nplt.tight_layout()\nplt.savefig(\"figure10.pdf\", dpi=300, bbox_inches=\"tight\", pad_inches=0.01)\n\n\n# norm_perf = []\n# norm_perf_ttft = []\n# for input_length in [256, 512, 1024, 2048]:\n# temp_list = []\n# our_prefill_latency = our_prefill[our_prefill[\"s\"] == input_length][\n# \"latency\"\n# ].values[0]\n# A100_prefill_latency = A100_prefill[A100_prefill[\"s\"] == input_length][\n# \"latency\"\n# ].values[0]\n# for output_length in [256, 512, 768, 1024, 1280, 1536, 1792, 2048]:\n# our_tbt_latency = get_total_decoding_latency(\n# our_decoding, input_length, input_length + output_length\n# )\n# A100_tbt_latency = get_total_decoding_latency(\n# A100_decoding, input_length, input_length + output_length\n# )\n# temp_list.append(our_tbt_latency / A100_tbt_latency)\n# norm_perf.append(temp_list)\n# norm_perf_ttft.append(our_prefill_latency / A100_prefill_latency)\n\n# cmap = sns.color_palette(\"viridis\", as_cmap=True)\n# data = np.array(norm_perf)\n# data_ttft = np.array(norm_perf_ttft)\n# print(data)\n# print(data_ttft)\n# import statistics\n# from matplotlib import gridspec\n\n# print(statistics.geometric_mean(data.flatten()))\n# print(statistics.geometric_mean(data_ttft))\n# # fig, axs = plt.subplots(1, 2, figsize=(8, 4),\n# # \t\t\tgridspec_kw={'width_ratios': [3, 1]}, sharey=True)\n# fig = plt.figure(figsize=(8, 3)) # Define the figure size\n# gs = gridspec.GridSpec(\n# 1, 2, width_ratios=[4, 1]\n# ) # 2 rows, 1 column, with the first row 3 times the height of the second\n# ax = fig.add_subplot(gs[0])\n# # ax=axs[0]\n# cax = ax.imshow(data, interpolation=\"nearest\", cmap=cmap, vmin=1.015, vmax=1.045)\n# # cax = sns.heatmap(data, cmap=\"viridis\")\n\n# # Add a colorbar\n# fig.colorbar(cax, shrink=1)\n\n\n# # Function to convert RGB to grayscale intensity\n# def get_intensity(color):\n# return color[0] * 0.299 + color[1] * 0.587 + color[2] * 0.114\n\n\n# # Set a threshold for deciding text color\n# intensity_threshold = 0.5\n# for i in range(data.shape[0]):\n# for j in range(data.shape[1]):\n# # Get the color from the colormap\n# cell_color = cax.cmap(cax.norm(data[i, j]))\n# # Calculate intensity of the cell color\n# intensity = get_intensity(cell_color)\n# # Choose text color based on intensity\n# text_color = \"white\" if intensity < intensity_threshold else \"black\"\n# text = ax.text(\n# j, i, round(data[i, j], 3), ha=\"center\", va=\"center\", color=text_color\n# )\n\n# # Set the x-axis and y-axis values\n# x_axis_labels = [256, 512, 768, 1024, 1280, 1536, 1792, 2048]\n# y_axis_labels = [256, 512, 1024, 2048]\n\n# # Set ticks positions\n# ax.set_xticks(np.arange(len(x_axis_labels)))\n# ax.set_yticks(np.arange(len(y_axis_labels)))\n\n\n# # Set ticks labels\n# ax.set_xticklabels(x_axis_labels)\n# ax.set_yticklabels(y_axis_labels)\n\n\n# # Set labels for axes\n# ax.set_xlabel(\"Output Length\\n\" + r\"$\\mathbf{Normalized\\ TBT}$\")\n# ax.set_ylabel(\"Input Length\")\n# ax.invert_yaxis()\n# # # Rotate the tick labels for the x-axis if needed\n# # plt.setp(ax.get_xticklabels(), rotation=45, ha=\"right\", rotation_mode=\"anchor\")\n\n\n# # fig = plt.figure(figsize=(10, 5)) # Define the figure size\n# axs1 = fig.add_subplot(gs[0, 1])\n# axs1.barh(np.arange(len(data_ttft)) / 2 + 0.2, data_ttft, color=\"steelblue\", height=0.3)\n# axs1.set_yticks(np.arange(len(y_axis_labels)) / 2 + 0.2)\n# axs1.set_yticklabels(y_axis_labels)\n# axs1.set_xlabel(r\"$\\mathbf{Normalized\\ TTFT}$\")\n# axs1.set_xlim(1, 2)\n\n# # Show the plot\n# plt.tight_layout()\n# plt.savefig(\"figure11.pdf\", dpi=300, bbox_inches=\"tight\", pad_inches=0.01)\n" }, { "path": "ae/figure10/run_figure10.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../..\n\npython -m ae.figure10.test_latency\n\ncd ae/figure10\npython plot_latency.py" }, { "path": "ae/figure10/test_latency.py", "content": "from software_model.transformer import (\n TransformerBlockInitComputationTP,\n TransformerBlockAutoRegressionTP,\n)\nfrom software_model.utils import data_type_dict, Tensor\nfrom design_space_exploration.dse import template_to_system, read_architecture_template\nfrom multiprocessing import Process, Lock\nfrom cost_model.cost_model import calc_compute_chiplet_area_mm2, calc_io_die_area_mm2\nimport time\n\nA100_specs = read_architecture_template(\"configs/GA100.json\")\nA100_system = template_to_system(A100_specs)\nour_specs = read_architecture_template(\"configs/latency_design.json\")\nour_system = template_to_system(our_specs)\nA100_compute_area_mm2 = calc_compute_chiplet_area_mm2(A100_specs)\nA100_io_area_mm2 = calc_io_die_area_mm2(A100_specs)\nour_compute_area_mm2 = calc_compute_chiplet_area_mm2(our_specs)\nour_io_area_mm2 = calc_io_die_area_mm2(our_specs)\nprint(f\"A100 compute area: {A100_compute_area_mm2} mm2\")\nprint(f\"A100 IO area: {A100_io_area_mm2} mm2\")\nprint(f\"A100 total area: {A100_compute_area_mm2+A100_io_area_mm2} mm2\")\nprint(f\"Our compute area: {our_compute_area_mm2} mm2\")\nprint(f\"Our IO area: {our_io_area_mm2} mm2\")\nprint(f\"Our total area: {our_compute_area_mm2+our_io_area_mm2} mm2\")\n\nwith open(\"ae/figure10/area.csv\", \"w\") as f:\n f.write(f\"A100 compute area: {A100_compute_area_mm2} mm2\\n\")\n f.write(f\"A100 IO area: {A100_io_area_mm2} mm2\\n\")\n f.write(f\"A100 total area: {A100_compute_area_mm2+A100_io_area_mm2} mm2\\n\")\n f.write(f\"Our compute area: {our_compute_area_mm2} mm2\\n\")\n f.write(f\"Our IO area: {our_io_area_mm2} mm2\\n\")\n f.write(f\"Our total area: {our_compute_area_mm2+our_io_area_mm2} mm2\\n\")\n\n\ndef simulate_decoding_latency(system, bs, seq_len, name, lock):\n model_auto_regression = TransformerBlockAutoRegressionTP(\n d_model=12288,\n n_heads=96,\n device_count=4,\n data_type=data_type_dict[\"fp16\"],\n )\n _ = model_auto_regression(\n Tensor([bs, 1, 12288], data_type_dict[\"fp16\"]),\n seq_len,\n )\n auto_regression_latency_simulated = model_auto_regression.compile_and_simulate(\n system, \"heuristic-GPU\"\n )\n with lock:\n with open(f\"ae/figure10/{name}_decoding.csv\", \"a\") as f:\n f.write(f\"{bs}, {seq_len}, {auto_regression_latency_simulated}\\n\")\n\n\ndef simulate_prefill_latency(system, bs, seq_len, name, lock):\n model = TransformerBlockInitComputationTP(\n d_model=12288,\n n_heads=96,\n device_count=4,\n data_type=data_type_dict[\"fp16\"],\n )\n _ = model(\n Tensor([bs, seq_len, 12288], data_type_dict[\"fp16\"]),\n )\n latency_simulated = model.compile_and_simulate(system, \"heuristic-GPU\")\n with lock:\n with open(f\"ae/figure10/{name}_prefill.csv\", \"a\") as f:\n f.write(f\"{bs}, {seq_len}, {latency_simulated}\\n\")\n\n\nlock_our_prefill = Lock()\nlock_our_decoding = Lock()\nlock_A100_prefill = Lock()\nlock_A100_decoding = Lock()\n\nprocesses = []\nfor bs in [16]: # [1, 4, 8, 16, 32, 64]:\n for seq_len in [256, 512, 1024, 2048]:\n for system in [our_system, A100_system]:\n if system == A100_system:\n name = \"A100\"\n lock = lock_A100_prefill\n else:\n name = \"our\"\n lock = lock_our_prefill\n p = Process(\n target=simulate_prefill_latency, args=(system, bs, seq_len, name, lock)\n )\n processes.append(p)\n for seq_len in range(256, 4096 + 64, 64):\n for system in [our_system, A100_system]:\n if system == A100_system:\n name = \"A100\"\n lock = lock_A100_decoding\n else:\n name = \"our\"\n lock = lock_our_decoding\n p = Process(\n target=simulate_decoding_latency, args=(system, bs, seq_len, name, lock)\n )\n processes.append(p)\n\n\ntry:\n for p in processes:\n p.start()\n print(\"Processes started.\")\n print(\"number of process:\", len(processes))\n while any(p.is_alive() for p in processes):\n time.sleep(1)\nexcept KeyboardInterrupt:\n print(\"Terminating processes...\")\n for p in processes:\n p.terminate()\n p.join()\n\n\nprint(\"All processes have finished.\")\n" }, { "path": "ae/figure11/__init__.py", "content": "" }, { "path": "ae/figure11/plot_decoding.py", "content": "import pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport numpy as np\nfrom scipy.stats import gmean\n\ncategories = [\"bs\", \"seq_len\", \"latency\"] + [\n \"Q_K_V\",\n \"Q_mul_K\",\n \"A_mul_V\",\n \"Wo_proj\",\n \"W1_proj\",\n \"W2_proj\",\n \"Softmax\",\n \"LayerNorm_MHA\",\n \"LayerNorm_FFN\",\n \"GeLU\",\n \"AllReduce_MHA\",\n \"AllReduce_FFN\",\n]\n\nA100 = pd.read_csv(\"A100.csv\", header=None, names=categories)\nA100[\"latency\"] = A100[\"latency\"] * 1000\nour = pd.read_csv(\"our.csv\", header=None, names=categories)\nour[\"latency\"] = our[\"latency\"] * 1000\n\nbs_list = [1, 2, 4, 8, 16, 32]\ncolors_our = sns.color_palette(\"Blues\", 3)[1:]\ncolors_a100 = sns.color_palette(\"summer_r\", 2)\nour_512 = our[our.seq_len == 512][our[\"bs\"].isin(bs_list)][\"latency\"].tolist()\nour_2048 = our[our.seq_len == 2048][our[\"bs\"].isin(bs_list)][\"latency\"].tolist()\na100_512 = A100[A100.seq_len == 512][A100[\"bs\"].isin(bs_list)][\"latency\"].tolist()\na100_2048 = A100[A100.seq_len == 2048][A100[\"bs\"].isin(bs_list)][\"latency\"].tolist()\navg_speedup = gmean(\n np.concatenate(\n (\n np.array(a100_512) / np.array(our_512),\n np.array(a100_2048) / np.array(our_2048),\n )\n )\n)\nprint(avg_speedup)\n\nplt.figure(figsize=(8, 3.5))\n\nx_pos = 0.25\nfor bs in bs_list:\n if bs == 1:\n seq_len = 512\n plt.bar(\n x_pos,\n A100[(A100.bs == bs) & (A100.seq_len == seq_len)].latency,\n width=0.5,\n label=f\"GA100 (seq_len={seq_len})\",\n color=colors_a100[0],\n )\n bars = plt.bar(\n x_pos + 0.5,\n our[(our.bs == bs) & (our.seq_len == seq_len)].latency,\n width=0.5,\n label=f\"Latency design (seq_len={seq_len})\",\n color=colors_our[0],\n )\n for bar in bars:\n bar.set_hatch(\"//\") # Add diagonal stripes\n seq_len = 2048\n plt.bar(\n x_pos + 1,\n A100[(A100.bs == bs) & (A100.seq_len == seq_len)].latency,\n width=0.5,\n label=f\"GA100 (seq_len={seq_len})\",\n color=colors_a100[1],\n )\n bars = plt.bar(\n x_pos + 1.5,\n our[(our.bs == bs) & (our.seq_len == seq_len)].latency,\n width=0.5,\n label=f\"Latency design (seq_len={seq_len})\",\n color=colors_our[1],\n )\n for bar in bars:\n bar.set_hatch(\"//\") # Add diagonal stripes\n else:\n seq_len = 512\n plt.bar(\n x_pos,\n A100[(A100.bs == bs) & (A100.seq_len == seq_len)].latency,\n width=0.5,\n color=colors_a100[0],\n )\n bars = plt.bar(\n x_pos + 0.5,\n our[(our.bs == bs) & (our.seq_len == seq_len)].latency,\n width=0.5,\n color=colors_our[0],\n )\n for bar in bars:\n bar.set_hatch(\"//\") # Add diagonal stripes\n seq_len = 2048\n if bs < 164:\n plt.bar(\n x_pos + 1,\n A100[(A100.bs == bs) & (A100.seq_len == seq_len)].latency,\n width=0.5,\n color=colors_a100[1],\n )\n bars = plt.bar(\n x_pos + 1.5,\n our[(our.bs == bs) & (our.seq_len == seq_len)].latency,\n width=0.5,\n color=colors_our[1],\n )\n for bar in bars:\n bar.set_hatch(\"//\") # Add diagonal stripes\n x_pos += 3\n\nplt.xticks([1, 4, 7, 10, 13, 16], bs_list)\nplt.xlabel(\"Batch Size\")\nplt.ylabel(\"Latency (ms)\")\nplt.legend(loc=\"upper left\")\nplt.tight_layout()\nplt.grid(True, axis=\"y\", ls=\"--\", c=\"0.8\")\nplt.savefig(\"figure11.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)\nplt.show()" }, { "path": "ae/figure11/run_figure11.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../..\n\npython -m ae.figure11.test_decoding\n\ncd ae/figure11\npython plot_decoding.py" }, { "path": "ae/figure11/test_decoding.py", "content": "from software_model.transformer import (\n TransformerBlockInitComputationTP,\n TransformerBlockAutoRegressionTP,\n)\nfrom software_model.utils import data_type_dict, Tensor\nfrom design_space_exploration.dse import template_to_system, read_architecture_template\nfrom multiprocessing import Process, Lock\nfrom cost_model.cost_model import calc_compute_chiplet_area_mm2, calc_io_die_area_mm2\nimport time\n\nA100_specs = read_architecture_template(\"configs/GA100.json\")\nA100_system = template_to_system(A100_specs)\nour_specs = read_architecture_template(\"configs/latency_design.json\")\nour_system = template_to_system(our_specs)\nA100_compute_area_mm2 = calc_compute_chiplet_area_mm2(A100_specs)\nA100_io_area_mm2 = calc_io_die_area_mm2(A100_specs)\nour_compute_area_mm2 = calc_compute_chiplet_area_mm2(our_specs)\nour_io_area_mm2 = calc_io_die_area_mm2(our_specs)\nprint(f\"A100 compute area: {A100_compute_area_mm2} mm2\")\nprint(f\"A100 IO area: {A100_io_area_mm2} mm2\")\nprint(f\"A100 total area: {A100_compute_area_mm2+A100_io_area_mm2} mm2\")\nprint(f\"Our compute area: {our_compute_area_mm2} mm2\")\nprint(f\"Our IO area: {our_io_area_mm2} mm2\")\nprint(f\"Our total area: {our_compute_area_mm2+our_io_area_mm2} mm2\")\n\n\ndef simulate_latency(system, bs, seq_len, name, lock):\n model_auto_regression = TransformerBlockAutoRegressionTP(\n d_model=12288,\n n_heads=96,\n device_count=4,\n data_type=data_type_dict[\"fp16\"],\n )\n _ = model_auto_regression(\n Tensor([bs, 1, 12288], data_type_dict[\"fp16\"]),\n seq_len,\n )\n auto_regression_latency_simulated = model_auto_regression.compile_and_simulate(\n system, \"heuristic-GPU\"\n )\n with lock:\n with open(f\"ae/figure11/{name}.csv\", \"a\") as f:\n f.write(\n f\"{bs}, {seq_len}, {auto_regression_latency_simulated}, {model_auto_regression.simluate_log}\\n\"\n )\n\n\nlock = Lock()\n\nprocesses = []\nfor bs in [1, 2, 4, 8, 16, 32, 64]:\n for seq_len in [512, 2048]:\n for system in [our_system, A100_system]:\n if system == A100_system:\n name = \"A100\"\n else:\n name = \"our\"\n p = Process(target=simulate_latency, args=(system, bs, seq_len, name, lock))\n processes.append(p)\n\ntry:\n for p in processes:\n p.start()\n\n while any(p.is_alive() for p in processes):\n time.sleep(1)\nexcept KeyboardInterrupt:\n print(\"Terminating processes...\")\n for p in processes:\n p.terminate()\n p.join()\n\n\nprint(\"All processes have finished.\")\n" }, { "path": "ae/figure12/__init__.py", "content": "" }, { "path": "ae/figure12/plot_throughput.py", "content": "import pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport statistics\n\nour_directory = \"our/our\"\nA100_directory = \"A100/A100\"\n\ncategories = [\"bs\", \"s\", \"latency\"] + [\n \"Q_K_V\",\n \"Q_mul_K\",\n \"A_mul_V\",\n \"Wo_proj\",\n \"W1_proj\",\n \"W2_proj\",\n \"Softmax\",\n \"LayerNorm_MHA\",\n \"LayerNorm_FFN\",\n \"GeLU\",\n \"AllReduce_MHA\",\n \"AllReduce_FFN\",\n]\n\nthroughput_our = []\nbs_our = []\nlatency_our = []\nthroughput_A100 = []\nbs_A100 = []\nlatency_A100 = []\n\n\ndef get_total_decoding_latency(df: pd.DataFrame, start, end):\n df_filtered = df[(df[\"s\"] >= start) & (df[\"s\"] <= end)]\n total_latency = 0\n\n # Calculate the mean of the values for each length interval and add to the sum\n for i in range(len(df_filtered) - 1):\n # Calculate the mean of current and next value\n mean = (df_filtered.iloc[i][\"latency\"] + df_filtered.iloc[i + 1][\"latency\"]) / 2\n # Calculate the difference in length\n length_interval = df_filtered.iloc[i + 1][\"s\"] - df_filtered.iloc[i][\"s\"]\n # Multiply the mean value by the length interval and add to the sum\n total_latency += mean * length_interval\n\n # print(total_latency)\n return total_latency\n\n\nfor input_length in [256, 512, 1024, 2048]:\n temp_our = []\n temp_A100 = []\n temp_our_bs = []\n temp_A100_bs = []\n temp_our_latency = []\n temp_A100_latency = []\n for output_length in [256, 512, 768, 1024, 1280, 1536, 1792, 2048]:\n our_prefill_df = pd.read_csv(\n f\"{our_directory}_{input_length}_{output_length}_prefill.csv\",\n header=None,\n names=categories,\n )\n # print(our_prefill_df)\n our_prefill_latency = our_prefill_df.iloc[0][\"latency\"]\n our_bs = our_prefill_df.iloc[0][\"bs\"]\n temp_our_bs.append(our_bs)\n our_decoding_df = pd.read_csv(\n f\"{our_directory}_{input_length}_{output_length}_decoding.csv\",\n header=None,\n names=categories,\n ).sort_values(by=\"s\")\n our_decoding_latency = get_total_decoding_latency(\n our_decoding_df, input_length, input_length + output_length\n )\n # print(our_decoding_latency)\n our_throughput = (\n our_bs * output_length / (our_prefill_latency + our_decoding_latency) / 12\n )\n temp_our.append(our_throughput)\n temp_our_latency.append(our_prefill_latency + our_decoding_latency)\n\n A100_prefill_df = pd.read_csv(\n f\"{A100_directory}_{input_length}_{output_length}_prefill.csv\",\n header=None,\n names=categories,\n )\n A100_prefill_latency = A100_prefill_df.iloc[0][\"latency\"]\n A100_bs = A100_prefill_df.iloc[0][\"bs\"]\n temp_A100_bs.append(A100_bs)\n A100_decoding_df = pd.read_csv(\n f\"{A100_directory}_{input_length}_{output_length}_decoding.csv\",\n header=None,\n names=categories,\n ).sort_values(by=\"s\")\n A100_decoding_latency = get_total_decoding_latency(\n A100_decoding_df, input_length, input_length + output_length\n )\n A100_throughput = (\n A100_bs\n * output_length\n / (A100_prefill_latency + A100_decoding_latency)\n / 12\n )\n temp_A100.append(A100_throughput)\n temp_A100_latency.append(A100_prefill_latency + A100_decoding_latency)\n\n throughput_our.append(temp_our)\n throughput_A100.append(temp_A100)\n bs_our.append(temp_our_bs)\n bs_A100.append(temp_A100_bs)\n latency_our.append(temp_our_latency)\n latency_A100.append(temp_A100_latency)\n# print(throughput_our)\n# print(throughput_A100)\nprint(latency_our)\nprint(latency_A100)\nprint(\n statistics.geometric_mean(\n (np.array(latency_our) / np.array(latency_A100)).flatten()\n )\n)\n\n\n# Function to convert RGB to grayscale intensity\ndef get_intensity(color):\n return color[0] * 0.299 + color[1] * 0.587 + color[2] * 0.114\n\n\ncmap = sns.color_palette(\"viridis\", as_cmap=True)\ndata = np.array(throughput_our) # / np.array(throughput_A100)\nprint(data.mean())\nfig, ax = plt.subplots()\ncax = ax.imshow(data, interpolation=\"nearest\", cmap=cmap)\n# cax = sns.heatmap(data, cmap=\"Blues\")\n\n# Add a colorbar\nfig.colorbar(cax, shrink=0.5)\n\n# Set a threshold for deciding text color\nintensity_threshold = 0.5\nfor i in range(data.shape[0]):\n for j in range(data.shape[1]):\n # Get the color from the colormap\n cell_color = cax.cmap(cax.norm(data[i, j]))\n # Calculate intensity of the cell color\n intensity = get_intensity(cell_color)\n # Choose text color based on intensity\n text_color = \"white\" if intensity < intensity_threshold else \"black\"\n text = ax.text(\n j, i, int(data[i, j]), ha=\"center\", va=\"center\", color=text_color\n )\n\n# Set the x-axis and y-axis values\nx_axis_labels = [256, 512, 768, 1024, 1280, 1536, 1792, 2048]\ny_axis_labels = [256, 512, 1024, 2048]\n\n# Set ticks positions\nax.set_xticks(np.arange(len(x_axis_labels)))\nax.set_yticks(np.arange(len(y_axis_labels)))\n\n# Set ticks labels\nax.set_xticklabels(x_axis_labels)\nax.set_yticklabels(y_axis_labels)\n\n# Set labels for axes\nax.set_xlabel(\"Output Length\")\nax.set_ylabel(\"Input Length\")\nax.invert_yaxis()\n# # Rotate the tick labels for the x-axis if needed\n# plt.setp(ax.get_xticklabels(), rotation=45, ha=\"right\", rotation_mode=\"anchor\")\n\n# Show the plot\nplt.tight_layout()\nplt.savefig(\"figure12a.pdf\", dpi=300, bbox_inches=\"tight\", pad_inches=0.01)\n\n\nimport statistics\n\ncmap = sns.color_palette(\"viridis\", as_cmap=True)\ndata = np.array(throughput_our) / np.array(throughput_A100)\nprint(statistics.geometric_mean(data.flatten()))\nfig, ax = plt.subplots()\ncax = ax.imshow(\n data,\n interpolation=\"nearest\",\n cmap=cmap,\n)\n# cax = sns.heatmap(data, cmap=\"viridis\")\n\n# Add a colorbar\nfig.colorbar(cax, shrink=0.5)\n\n# Set a threshold for deciding text color\nintensity_threshold = 0.5\nfor i in range(data.shape[0]):\n for j in range(data.shape[1]):\n # Get the color from the colormap\n cell_color = cax.cmap(cax.norm(data[i, j]))\n # Calculate intensity of the cell color\n intensity = get_intensity(cell_color)\n # Choose text color based on intensity\n text_color = \"white\" if intensity < intensity_threshold else \"black\"\n text = ax.text(\n j, i, round(data[i, j], 2), ha=\"center\", va=\"center\", color=text_color\n )\n\n# Set the x-axis and y-axis values\nx_axis_labels = [256, 512, 768, 1024, 1280, 1536, 1792, 2048]\ny_axis_labels = [256, 512, 1024, 2048]\n\n# Set ticks positions\nax.set_xticks(np.arange(len(x_axis_labels)))\nax.set_yticks(np.arange(len(y_axis_labels)))\n\n# Set ticks labels\nax.set_xticklabels(x_axis_labels)\nax.set_yticklabels(y_axis_labels)\n\n# Set labels for axes\nax.set_xlabel(\"Output Length\")\nax.set_ylabel(\"Input Length\")\nax.invert_yaxis()\n# # Rotate the tick labels for the x-axis if needed\n# plt.setp(ax.get_xticklabels(), rotation=45, ha=\"right\", rotation_mode=\"anchor\")\n\n# Show the plot\nplt.tight_layout()\nplt.savefig(\"figure12b.pdf\", dpi=300, bbox_inches=\"tight\", pad_inches=0.01)\n" }, { "path": "ae/figure12/run_figure12.sh", "content": "rm A100/*.csv\nrm our/*.csv\nrm *.pdf\n\nmkdir A100\nmkdir our\n\ncd ../..\n\npython -m ae.figure12.test_throughput\n\ncd ae/figure12\npython plot_throughput.py" }, { "path": "ae/figure12/test_throughput.py", "content": "from software_model.transformer import (\n TransformerBlockInitComputationTP,\n TransformerBlockAutoRegressionTP,\n)\nfrom software_model.utils import data_type_dict, Tensor\nfrom design_space_exploration.dse import template_to_system, read_architecture_template\nfrom multiprocessing import Process, Lock\nfrom cost_model.cost_model import calc_compute_chiplet_area_mm2, calc_io_die_area_mm2\nimport time\n\nA100_specs = read_architecture_template(\"configs/GA100.json\")\nA100_system = template_to_system(A100_specs)\nour_specs = read_architecture_template(\"configs/prefilling_system.json\")\nour_system = template_to_system(our_specs)\nA100_compute_area_mm2 = calc_compute_chiplet_area_mm2(A100_specs)\nA100_io_area_mm2 = calc_io_die_area_mm2(A100_specs)\nour_compute_area_mm2 = calc_compute_chiplet_area_mm2(our_specs)\nour_io_area_mm2 = calc_io_die_area_mm2(our_specs)\nprint(f\"A100 compute area: {A100_compute_area_mm2} mm2\")\nprint(f\"A100 IO area: {A100_io_area_mm2} mm2\")\nprint(f\"A100 total area: {A100_compute_area_mm2+A100_io_area_mm2} mm2\")\nprint(f\"Our compute area: {our_compute_area_mm2} mm2\")\nprint(f\"Our IO area: {our_io_area_mm2} mm2\")\nprint(f\"Our total area: {our_compute_area_mm2+our_io_area_mm2} mm2\")\nwith open(\"ae/figure12/area.csv\", \"w\") as f:\n f.write(f\"A100 compute area: {A100_compute_area_mm2} mm2\\n\")\n f.write(f\"A100 IO area: {A100_io_area_mm2} mm2\\n\")\n f.write(f\"A100 total area: {A100_compute_area_mm2+A100_io_area_mm2} mm2\\n\")\n f.write(f\"Our compute area: {our_compute_area_mm2} mm2\\n\")\n f.write(f\"Our IO area: {our_io_area_mm2} mm2\\n\")\n f.write(f\"Our total area: {our_compute_area_mm2+our_io_area_mm2} mm2\\n\")\n\n\ndef simulate_decoding_latency(system, bs, seq_len, name, lock, heuristics):\n model_auto_regression = TransformerBlockAutoRegressionTP(\n d_model=12288,\n n_heads=96,\n device_count=1,\n data_type=data_type_dict[\"fp16\"],\n )\n _ = model_auto_regression(\n Tensor([bs, 1, 12288], data_type_dict[\"fp16\"]),\n seq_len,\n )\n auto_regression_latency_simulated = model_auto_regression.compile_and_simulate(\n system, heuristics\n )\n with lock:\n with open(f\"ae/figure12/{name}_decoding.csv\", \"a\") as f:\n f.write(\n f\"{bs}, {seq_len}, {auto_regression_latency_simulated}, {model_auto_regression.simluate_log}\\n\"\n )\n\n\ndef simulate_prefill_latency(system, bs, seq_len, name, lock, heuristics):\n model = TransformerBlockInitComputationTP(\n d_model=12288,\n n_heads=96,\n device_count=1,\n data_type=data_type_dict[\"fp16\"],\n )\n _ = model(\n Tensor([bs, seq_len, 12288], data_type_dict[\"fp16\"]),\n )\n latency_simulated = model.compile_and_simulate(system, heuristics)\n with lock:\n with open(f\"ae/figure12/{name}_prefill.csv\", \"a\") as f:\n f.write(f\"{bs}, {seq_len}, {latency_simulated}, {model.simluate_log}\\n\")\n\n\nlock_our_prefill = Lock()\nlock_our_decoding = Lock()\nlock_A100_prefill = Lock()\nlock_A100_decoding = Lock()\n\nprocesses = []\n\nfor input_seq_len in [\n 256,\n 512,\n 1024,\n 2048,\n]:\n for output_seq_len in [\n 256,\n 512,\n 768,\n 1024,\n 1280,\n 1536,\n 1792,\n 2048,\n ]:\n seq_len = input_seq_len + output_seq_len\n for system in [our_system, A100_system]:\n if system == A100_system:\n name = f\"A100/A100_{input_seq_len}_{output_seq_len}\"\n lock = lock_A100_prefill\n bs = (80e9 - 2 * 12 * 12288**2 * 12) // ((12 * 4 + 8) * seq_len * 12288)\n heuristics = \"heuristic-GPU\"\n else:\n name = f\"our/our_{input_seq_len}_{output_seq_len}\"\n lock = lock_our_prefill\n bs = (512e9 - 2 * 12 * 12288**2 * 12) // (\n (12 * 4 + 8) * seq_len * 12288\n )\n heuristics = \"heuristic-our-throughput\"\n bs = int(bs)\n # print(bs)\n p = Process(\n target=simulate_prefill_latency,\n args=(system, bs, input_seq_len, name, lock, heuristics),\n )\n processes.append(p)\n\n for decoding_seq_len in range(input_seq_len, seq_len + 64, 64):\n if system == A100_system:\n name = f\"A100/A100_{input_seq_len}_{output_seq_len}\"\n lock = lock_A100_decoding\n heuristics = \"heuristic-GPU\"\n else:\n name = f\"our/our_{input_seq_len}_{output_seq_len}\"\n lock = lock_our_decoding\n heuristics = \"heuristic-our-throughput\"\n p = Process(\n target=simulate_decoding_latency,\n args=(system, bs, decoding_seq_len, name, lock, heuristics),\n )\n processes.append(p)\n\nprint(len(processes))\n# exit()\ntry:\n for p in processes:\n p.start()\n print(\"Processes started.\")\n print(\"number of process:\", len(processes))\n while any(p.is_alive() for p in processes):\n time.sleep(1)\nexcept KeyboardInterrupt:\n print(\"Terminating processes...\")\n for p in processes:\n p.terminate()\n p.join()\n\n\nprint(\"All processes have finished.\")\n" }, { "path": "ae/figure5/__init__.py", "content": "" }, { "path": "ae/figure5/ab/__init__.py", "content": "" }, { "path": "ae/figure5/ab/plot_matmul.py", "content": "import pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport numpy as np\n\n\nmatmul_TPUv3_sim = pd.read_csv(\n \"matmul_TPUv3_sim.csv\", header=None, names=[\"M\", \"N\", \"K\", \"latency\", \"throughput\"]\n)\nmatmul_TPUv3_sim[\"throughput\"] = (\n matmul_TPUv3_sim[\"throughput\"].str.extract(r\"(\\d+\\.?\\d*)\").astype(float)\n)\nmatmul_TPUv3_sim.set_index([\"M\", \"N\", \"K\"], inplace=True)\nmatmul_TPUv3_roofline = pd.read_csv(\n \"matmul_TPUv3_roofline.csv\",\n header=None,\n names=[\"M\", \"N\", \"K\", \"latency\", \"throughput\"],\n)\nmatmul_TPUv3_roofline[\"throughput\"] = (\n matmul_TPUv3_roofline[\"throughput\"].str.extract(r\"(\\d+\\.?\\d*)\").astype(float)\n)\nmatmul_TPUv3_roofline.set_index([\"M\", \"N\", \"K\"], inplace=True)\nmatmul_A100 = pd.read_csv(\n \"real_hardware/matmul_A100.csv\", header=None, names=[\"M\", \"N\", \"K\", \"latency\", \"throughput\"]\n)\nmatmul_A100[\"throughput\"] = (\n matmul_A100[\"throughput\"].str.extract(r\"(\\d+\\.?\\d*)\").astype(float)\n)\nmatmul_A100.set_index([\"M\", \"N\", \"K\"], inplace=True)\nmatmul_A100_sim = pd.read_csv(\n \"matmul_A100_sim.csv\", header=None, names=[\"M\", \"N\", \"K\", \"latency\", \"throughput\"]\n)\nmatmul_A100_sim[\"throughput\"] = (\n matmul_A100_sim[\"throughput\"].str.extract(r\"(\\d+\\.?\\d*)\").astype(float)\n)\nmatmul_A100_sim.set_index([\"M\", \"N\", \"K\"], inplace=True)\nmatmul_A100_roofline = pd.read_csv(\n \"matmul_A100_roofline.csv\",\n header=None,\n names=[\"M\", \"N\", \"K\", \"latency\", \"throughput\"],\n)\nmatmul_A100_roofline[\"throughput\"] = (\n matmul_A100_roofline[\"throughput\"].str.extract(r\"(\\d+\\.?\\d*)\").astype(float)\n)\nmatmul_A100_roofline.set_index([\"M\", \"N\", \"K\"], inplace=True)\nmatmul_MI210 = pd.read_csv(\n \"real_hardware/matmul_MI210.csv\", header=None, names=[\"M\", \"N\", \"K\", \"latency\", \"throughput\"]\n)\nmatmul_MI210[\"throughput\"] = (\n matmul_MI210[\"throughput\"].str.extract(r\"(\\d+\\.?\\d*)\").astype(float)\n)\nmatmul_MI210.set_index([\"M\", \"N\", \"K\"], inplace=True)\nmatmul_MI210_sim = pd.read_csv(\n \"matmul_MI210_sim.csv\", header=None, names=[\"M\", \"N\", \"K\", \"latency\", \"throughput\"]\n)\nmatmul_MI210_sim[\"throughput\"] = (\n matmul_MI210_sim[\"throughput\"].str.extract(r\"(\\d+\\.?\\d*)\").astype(float)\n)\nmatmul_MI210_sim.set_index([\"M\", \"N\", \"K\"], inplace=True)\nmatmul_MI210_roofline = pd.read_csv(\n \"matmul_MI210_roofline.csv\",\n header=None,\n names=[\"M\", \"N\", \"K\", \"latency\", \"throughput\"],\n)\nmatmul_MI210_roofline[\"throughput\"] = (\n matmul_MI210_roofline[\"throughput\"].str.extract(r\"(\\d+\\.?\\d*)\").astype(float)\n)\nmatmul_MI210_roofline.set_index([\"M\", \"N\", \"K\"], inplace=True)\n\ncolor_NV = sns.color_palette(\"Greens_d\", 4)[1:]\ncolor_Google = sns.color_palette(\"Blues_d\", 4)[1:]\ncolor_AMD = sns.color_palette(\"flare\", 3)\n\nK = 12288\nN = K\ntitle = f\"Performance of Matmul with K={K}, N={N}\"\nM_list = []\nthroughput_TPU_list = []\nthroughput_TPU_sim_list = []\nthroughput_TPU_roofline_list = []\nthroughput_GPU_list = []\nthroughput_GPU_sim_list = []\nthroughput_GPU_roofline_list = []\nthroughput_AMD_list = []\nthroughput_AMD_sim_list = []\nthroughput_AMD_roofline_list = []\nfor M in range(6, 16):\n M = 2**M\n M_list.append(M)\n throughput_TPU_sim_list.append(matmul_TPUv3_sim.loc[(M, N, K), \"throughput\"])\n throughput_TPU_roofline_list.append(\n matmul_TPUv3_roofline.loc[(M, N, K), \"throughput\"]\n )\n throughput_GPU_list.append(matmul_A100.loc[(M, N, K), \"throughput\"])\n throughput_GPU_sim_list.append(matmul_A100_sim.loc[(M, N, K), \"throughput\"])\n throughput_GPU_roofline_list.append(\n matmul_A100_roofline.loc[(M, N, K), \"throughput\"]\n )\n throughput_AMD_list.append(matmul_MI210.loc[(M, N, K), \"throughput\"])\n throughput_AMD_sim_list.append(matmul_MI210_sim.loc[(M, N, K), \"throughput\"])\n throughput_AMD_roofline_list.append(\n matmul_MI210_roofline.loc[(M, N, K), \"throughput\"]\n )\n\n\n# plt.figure(figsize=(6, 2.8))\nplt.figure(figsize=(3.64, 2.8))\nplt.xscale(\"log\", base=2)\nplt.plot(\n M_list,\n throughput_GPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of NVIDIA A100\",\n color=color_NV[0],\n)\nplt.plot(\n M_list, throughput_GPU_list, marker=\"o\", label=\"Real NVIDIA A100\", color=color_NV[1]\n)\nplt.plot(\n M_list,\n throughput_GPU_sim_list,\n marker=\"x\",\n label=\"Simulated NVIDIA A100\",\n color=color_NV[2],\n)\nplt.plot(\n M_list,\n throughput_AMD_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of AMD MI210\",\n color=color_AMD[0],\n)\nplt.plot(\n M_list, throughput_AMD_list, marker=\"o\", label=\"Real AMD MI210\", color=color_AMD[1]\n)\nplt.plot(\n M_list,\n throughput_AMD_sim_list,\n marker=\"x\",\n label=\"Simulated AMD MI210\",\n color=color_AMD[2],\n)\nplt.plot(\n M_list,\n throughput_TPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of Google TPUv3\",\n color=color_Google[0],\n)\nplt.plot(\n M_list,\n throughput_TPU_sim_list,\n marker=\"x\",\n label=\"Simulated Google TPUv3\",\n color=color_Google[2],\n)\n\n\n# handles, labels = plt.gca().get_legend_handles_labels()\n# plt.legend(handles, labels, loc=\"upper left\", bbox_to_anchor=(1, 1))\n# plt.title(title)\nplt.xlabel(\"M\")\nplt.ylabel(\"TFLOPS\")\nplt.grid(True, which=\"both\", ls=\"--\", c=\"0.7\") # Adding a grid for better readability\nplt.tight_layout()\nplt.savefig(\"figure5b.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)\n\n\nM = 8192\ntitle = f\"Performance of Matmul with M={M}\"\nK_list = []\nthroughput_TPU_list = []\nthroughput_TPU_sim_list = []\nthroughput_TPU_roofline_list = []\nthroughput_GPU_list = []\nthroughput_GPU_sim_list = []\nthroughput_GPU_roofline_list = []\nthroughput_AMD_list = []\nthroughput_AMD_sim_list = []\nthroughput_AMD_roofline_list = []\nfor K in range(6, 16):\n K = 2**K\n N = K\n K_list.append(K)\n throughput_TPU_sim_list.append(matmul_TPUv3_sim.loc[(M, N, K), \"throughput\"])\n throughput_TPU_roofline_list.append(\n matmul_TPUv3_roofline.loc[(M, N, K), \"throughput\"]\n )\n throughput_GPU_list.append(matmul_A100.loc[(M, N, K), \"throughput\"])\n throughput_GPU_sim_list.append(matmul_A100_sim.loc[(M, N, K), \"throughput\"])\n throughput_GPU_roofline_list.append(\n matmul_A100_roofline.loc[(M, N, K), \"throughput\"]\n )\n throughput_AMD_list.append(matmul_MI210.loc[(M, N, K), \"throughput\"])\n throughput_AMD_sim_list.append(matmul_MI210_sim.loc[(M, N, K), \"throughput\"])\n throughput_AMD_roofline_list.append(\n matmul_MI210_roofline.loc[(M, N, K), \"throughput\"]\n )\n\n\nplt.figure(figsize=(3.64, 2.8))\nplt.xscale(\"log\", base=2)\nplt.plot(\n K_list,\n throughput_GPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of NVIDIA A100\",\n color=color_NV[0],\n)\nplt.plot(\n K_list, throughput_GPU_list, marker=\"o\", label=\"Real NVIDIA A100\", color=color_NV[1]\n)\nplt.plot(\n K_list,\n throughput_GPU_sim_list,\n marker=\"x\",\n label=\"Simulated NVIDIA A100\",\n color=color_NV[2],\n)\nplt.plot(\n K_list,\n throughput_AMD_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of AMD MI210\",\n color=color_AMD[0],\n)\nplt.plot(\n K_list, throughput_AMD_list, marker=\"o\", label=\"Real AMD MI210\", color=color_AMD[1]\n)\nplt.plot(\n K_list,\n throughput_AMD_sim_list,\n marker=\"x\",\n label=\"Simulated AMD MI210\",\n color=color_AMD[2],\n)\nplt.plot(\n K_list,\n throughput_TPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of Google TPUv3\",\n color=color_Google[0],\n)\nplt.plot(\n K_list,\n throughput_TPU_sim_list,\n marker=\"x\",\n label=\"Simulated Google TPUv3\",\n color=color_Google[2],\n)\n# plt.legend()\n# plt.title(title)\nplt.xlabel(\"N=K\")\nplt.ylabel(\"TFLOPS\")\nplt.grid(True, which=\"both\", ls=\"--\", c=\"0.7\") # Adding a grid for better readability\nplt.tight_layout()\nplt.savefig(\"figure5a.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)" }, { "path": "ae/figure5/ab/real_hardware/matmul_A100.csv", "content": "64, 12288, 12288, 0.1900ms, 101.7124Tflops\n128, 12288, 12288, 0.2003ms, 193.0114Tflops\n256, 12288, 12288, 0.3185ms, 242.7090Tflops\n512, 12288, 12288, 0.6118ms, 252.7351Tflops\n1024, 12288, 12288, 1.1990ms, 257.9115Tflops\n2048, 12288, 12288, 2.3586ms, 262.2263Tflops\n4096, 12288, 12288, 4.4576ms, 277.4929Tflops\n8192, 12288, 12288, 8.6216ms, 286.9431Tflops\n16384, 12288, 12288, 17.0307ms, 290.5223Tflops\n32768, 12288, 12288, 35.4407ms, 279.2160Tflops\n8192, 64, 64, 0.0296ms, 2.2700Tflops\n8192, 128, 128, 0.0310ms, 8.6608Tflops\n8192, 256, 256, 0.0356ms, 30.1244Tflops\n8192, 512, 512, 0.0471ms, 91.2121Tflops\n8192, 1024, 1024, 0.0927ms, 185.2380Tflops\n8192, 2048, 2048, 0.2818ms, 243.8497Tflops\n8192, 4096, 4096, 1.0210ms, 269.2169Tflops\n8192, 8192, 8192, 3.9614ms, 277.5532Tflops\n8192, 16384, 16384, 15.0087ms, 293.0334Tflops\n8192, 32768, 32768, 61.1346ms, 287.7616Tflops" }, { "path": "ae/figure5/ab/real_hardware/matmul_MI210.csv", "content": "32, 12288, 12288, 0.5493ms, 17.5922Tflops\n64, 12288, 12288, 0.5584ms, 34.6135Tflops\n128, 12288, 12288, 0.5932ms, 65.1646Tflops\n256, 12288, 12288, 0.7699ms, 100.4209Tflops\n512, 12288, 12288, 1.4054ms, 110.0209Tflops\n1024, 12288, 12288, 2.7173ms, 113.8051Tflops\n2048, 12288, 12288, 5.3905ms, 114.7338Tflops\n4096, 12288, 12288, 10.4494ms, 118.3752Tflops\n8192, 12288, 12288, 20.7849ms, 119.0242Tflops\n16384, 12288, 12288, 41.1353ms, 120.2811Tflops\n32768, 12288, 12288, 81.4046ms, 121.5608Tflops\n8192, 32, 32, 0.0333ms, 0.5044Tflops\n8192, 64, 64, 0.0345ms, 1.9479Tflops\n8192, 128, 128, 0.0396ms, 6.7825Tflops\n8192, 256, 256, 0.0485ms, 22.1307Tflops\n8192, 512, 512, 0.0863ms, 49.7635Tflops\n8192, 1024, 1024, 0.1950ms, 88.0900Tflops\n8192, 2048, 2048, 0.5822ms, 118.0305Tflops\n8192, 4096, 4096, 2.2901ms, 120.0272Tflops\n8192, 8192, 8192, 9.4150ms, 116.7826Tflops\n8192, 16384, 16384, 36.7552ms, 119.6578Tflops\n8192, 32768, 32768, 146.2553ms, 120.2841Tflops" }, { "path": "ae/figure5/ab/run.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../../..\n\npython -m ae.figure5.ab.test_matmul --simgpu --roofline\npython -m ae.figure5.ab.test_matmul --simtpu --roofline\npython -m ae.figure5.ab.test_matmul --simamd --roofline\n\npython -m ae.figure5.ab.test_matmul --simgpu\npython -m ae.figure5.ab.test_matmul --simtpu\npython -m ae.figure5.ab.test_matmul --simamd\n\ncd ae/figure5/ab\npython plot_matmul.py" }, { "path": "ae/figure5/ab/test_matmul.py", "content": "from software_model.matmul import Matmul\nfrom software_model.utils import data_type_dict, Tensor\nfrom hardware_model.device import device_dict\nimport argparse\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--gpu\", action=\"store_true\", help=\"Enable GPU\")\n parser.add_argument(\"--simtpu\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--simtpu-new\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--simgpu\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--simamd\", action=\"store_true\", help=\"amd simulation\")\n parser.add_argument(\"--roofline\", action=\"store_true\", help=\"Roofline simulation\")\n args = parser.parse_args()\n\n if args.simtpu:\n pcb = device_dict[\"TPUv3\"]\n if args.simtpu_new:\n pcb = device_dict[\"TPUv3_new\"]\n if args.simgpu:\n pcb = device_dict[\"A100_80GB_fp16\"]\n\n MI210 = device_dict[\"MI210\"]\n amd_overhead = MI210.compute_module.overhead.softmax\n\n K = 12288\n N = K\n titile = f\"Performance of Matmul with K={K}, N={N}\"\n print(f\"{titile}\")\n\n test_overhead = True\n\n for M in range(5, 16):\n M = 2**M\n model = Matmul(data_type=data_type_dict[\"fp16\"])\n _ = model(\n Tensor([M, K]),\n Tensor([K, N]),\n )\n if args.gpu:\n if test_overhead:\n model.gpu_kernel_launch_overhead()\n test_overhead = False\n latency = model.run_on_gpu()\n if args.simtpu:\n if args.roofline:\n latency = model.roofline_model(pcb) + 110e-6\n file_name='matmul_TPUv3_roofline.csv'\n else:\n latency = (\n model.compile_and_simulate(pcb, compile_mode=\"heuristic-TPU\")\n + 110e-6\n )\n file_name='matmul_TPUv3_sim.csv'\n\n if args.simtpu_new:\n if args.roofline:\n latency = model.roofline_model(pcb) + 110e-6\n else:\n latency = (\n model.compile_and_simulate(pcb, compile_mode=\"heuristic-TPU-new\")\n + 110e-6\n )\n if args.simgpu:\n if args.roofline:\n latency = model.roofline_model(pcb) + 2.1e-5\n file_name='matmul_A100_roofline.csv'\n else:\n latency = (\n model.compile_and_simulate(pcb, compile_mode=\"heuristic-GPU\")\n + 2.1e-5\n )\n file_name='matmul_A100_sim.csv'\n if args.simamd:\n if args.roofline:\n latency = model.roofline_model(pcb_module=MI210) + amd_overhead\n file_name='matmul_MI210_roofline.csv'\n else:\n latency = (\n model.compile_and_simulate(\n pcb_module=MI210, compile_mode=\"heuristic-GPU\"\n )\n + amd_overhead\n )\n file_name='matmul_MI210_sim.csv'\n tflops = 2 * M * N * K / latency / 1e12\n print(f\"{M}, {N}, {K}, {latency*1e3:.4f}ms, {tflops:.4f}Tflops\", flush=True)\n with open(f'ae/figure5/ab/{file_name}', 'a') as f:\n f.write(f\"{M}, {N}, {K}, {latency*1e3:.4f}ms, {tflops:.4f}Tflops\\n\")\n\n M = 8192\n print(f\"Performance of Matmul with M={M}, N=K\")\n for K in range(5, 16):\n K = 2**K\n N = K\n model = Matmul(data_type=data_type_dict[\"fp16\"])\n _ = model(\n Tensor([M, K]),\n Tensor([K, N]),\n )\n if args.gpu:\n latency = model.run_on_gpu()\n if args.simtpu:\n if args.roofline:\n latency = model.roofline_model(pcb) + 110e-6\n else:\n latency = (\n model.compile_and_simulate(pcb, compile_mode=\"heuristic-TPU\")\n + 110e-6\n )\n if args.simtpu_new:\n if args.roofline:\n latency = model.roofline_model(pcb) + 110e-6\n else:\n latency = (\n model.compile_and_simulate(pcb, compile_mode=\"heuristic-TPU-new\")\n + 110e-6\n )\n if args.simgpu:\n if args.roofline:\n latency = model.roofline_model(pcb) + 2.1e-5\n else:\n latency = (\n model.compile_and_simulate(pcb, compile_mode=\"heuristic-GPU\")\n + 2.1e-5\n )\n if args.simamd:\n if args.roofline:\n latency = model.roofline_model(pcb_module=MI210) + amd_overhead\n else:\n latency = (\n model.compile_and_simulate(\n pcb_module=MI210, compile_mode=\"heuristic-GPU\"\n )\n + amd_overhead\n )\n tflops = 2 * M * N * K / latency / 1e12\n print(f\"{M}, {N}, {K}, {latency*1e3:.4f}ms, {tflops:.4f}Tflops\", flush=True)\n with open(f'ae/figure5/ab/{file_name}', 'a') as f:\n f.write(f\"{M}, {N}, {K}, {latency*1e3:.4f}ms, {tflops:.4f}Tflops\\n\")\n" }, { "path": "ae/figure5/cf/__init__.py", "content": "" }, { "path": "ae/figure5/cf/plot_softmax.py", "content": "import pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport numpy as np\n\n\nsoftmax_TPUv3_sim = pd.read_csv(\n \"softmax_TPUv3_sim.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nsoftmax_TPUv3_sim.set_index([\"M\", \"N\"], inplace=True)\nsoftmax_TPUv3_roofline = pd.read_csv(\n \"softmax_TPUv3_roofline.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nsoftmax_TPUv3_roofline.set_index([\"M\", \"N\"], inplace=True)\nsoftmax_A100 = pd.read_csv(\n \"real_hardware/softmax_A100.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nsoftmax_A100.set_index([\"M\", \"N\"], inplace=True)\nsoftmax_A100_sim = pd.read_csv(\n \"softmax_A100_sim.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nsoftmax_A100_sim.set_index([\"M\", \"N\"], inplace=True)\nsoftmax_A100_roofline = pd.read_csv(\n \"softmax_A100_roofline.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nsoftmax_A100_roofline.set_index([\"M\", \"N\"], inplace=True)\nsoftmax_MI210 = pd.read_csv(\n \"real_hardware/softmax_MI210.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nsoftmax_MI210.set_index([\"M\", \"N\"], inplace=True)\nsoftmax_MI210_sim = pd.read_csv(\n \"softmax_MI210_sim.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nsoftmax_MI210_sim.set_index([\"M\", \"N\"], inplace=True)\nsoftmax_MI210_roofline = pd.read_csv(\n \"softmax_MI210_roofline.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nsoftmax_MI210_roofline.set_index([\"M\", \"N\"], inplace=True)\n\ncolor_NV = sns.color_palette(\"Greens_d\", 4)[1:]\ncolor_Google = sns.color_palette(\"Blues_d\", 4)[1:]\ncolor_AMD = sns.color_palette(\"flare\", 3)\n\nM = 4096\ntitle = f\"Performance of softmax with M={M}\"\nN_list = []\nthroughput_TPU_list = []\nthroughput_TPU_sim_list = []\nthroughput_TPU_roofline_list = []\nthroughput_GPU_list = []\nthroughput_GPU_sim_list = []\nthroughput_GPU_roofline_list = []\nthroughput_AMD_list = []\nthroughput_AMD_sim_list = []\nthroughput_AMD_roofline_list = []\nfor N in range(6, 16):\n N = 2**N\n N_list.append(N)\n # print(M,N)\n # print(softmax_TPUv3.loc[(M, N), 'throughput'])\n\n throughput_TPU_sim_list.append(\n softmax_TPUv3_sim.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_TPU_roofline_list.append(\n softmax_TPUv3_roofline.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_GPU_list.append(softmax_A100.loc[(M, N), \"throughput\"].values[0])\n throughput_GPU_sim_list.append(softmax_A100_sim.loc[(M, N), \"throughput\"].values[0])\n throughput_GPU_roofline_list.append(\n softmax_A100_roofline.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_AMD_list.append(softmax_MI210.loc[(M, N), \"throughput\"].values[0])\n throughput_AMD_sim_list.append(\n softmax_MI210_sim.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_AMD_roofline_list.append(\n softmax_MI210_roofline.loc[(M, N), \"throughput\"].values[0]\n )\n\nplt.figure(figsize=(3.7, 2))\nplt.xscale(\"log\", base=2)\nplt.plot(\n N_list,\n throughput_GPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of NVIDIA A100\",\n color=color_NV[0],\n)\nplt.plot(\n N_list, throughput_GPU_list, marker=\"o\", label=\"Real NVIDIA A100\", color=color_NV[1]\n)\nplt.plot(\n N_list,\n throughput_GPU_sim_list,\n marker=\"x\",\n label=\"Simulated NVIDIA A100\",\n color=color_NV[2],\n)\nplt.plot(\n N_list,\n throughput_AMD_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of AMD MI210\",\n color=color_AMD[0],\n)\nplt.plot(\n N_list, throughput_AMD_list, marker=\"o\", label=\"Real AMD MI210\", color=color_AMD[1]\n)\nplt.plot(\n N_list,\n throughput_AMD_sim_list,\n marker=\"x\",\n label=\"Simulated AMD MI210\",\n color=color_AMD[2],\n)\nplt.plot(\n N_list,\n throughput_TPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of Google TPUv3\",\n color=color_Google[0],\n)\n\nplt.plot(\n N_list,\n throughput_TPU_sim_list,\n marker=\"x\",\n label=\"Simulated Google TPUv3\",\n color=color_Google[2],\n)\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles, labels, loc=\"upper left\", bbox_to_anchor=(1, 1.1))\n# plt.legend()\n# plt.title(title)\nplt.xlabel(\"N\")\nplt.ylabel(\"G Elements/s\")\nplt.grid(True, which=\"both\", ls=\"--\", c=\"0.7\") # Adding a grid for better readability\nplt.savefig(\"figure5f.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)\n\n\nN = 4096\ntitle = f\"Performance of softmax with N={N}\"\nM_list = []\nthroughput_TPU_list = []\nthroughput_TPU_sim_list = []\nthroughput_TPU_roofline_list = []\nthroughput_GPU_list = []\nthroughput_GPU_sim_list = []\nthroughput_GPU_roofline_list = []\nthroughput_AMD_list = []\nthroughput_AMD_sim_list = []\nthroughput_AMD_roofline_list = []\nfor M in range(6, 16):\n M = 2**M\n M_list.append(M)\n throughput_TPU_sim_list.append(\n softmax_TPUv3_sim.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_TPU_roofline_list.append(\n softmax_TPUv3_roofline.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_GPU_list.append(softmax_A100.loc[(M, N), \"throughput\"].values[0])\n throughput_GPU_sim_list.append(softmax_A100_sim.loc[(M, N), \"throughput\"].values[0])\n throughput_GPU_roofline_list.append(\n softmax_A100_roofline.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_AMD_list.append(softmax_MI210.loc[(M, N), \"throughput\"].values[0])\n throughput_AMD_sim_list.append(\n softmax_MI210_sim.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_AMD_roofline_list.append(\n softmax_MI210_roofline.loc[(M, N), \"throughput\"].values[0]\n )\n\n\nplt.figure(figsize=(3.7, 2.8))\nplt.xscale(\"log\", base=2)\nplt.plot(\n M_list,\n throughput_GPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of NVIDIA A100\",\n color=color_NV[0],\n)\nplt.plot(\n M_list, throughput_GPU_list, marker=\"o\", label=\"Real NVIDIA A100\", color=color_NV[1]\n)\nplt.plot(\n M_list,\n throughput_GPU_sim_list,\n marker=\"x\",\n label=\"Simulated NVIDIA A100\",\n color=color_NV[2],\n)\nplt.plot(\n M_list,\n throughput_AMD_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of AMD MI210\",\n color=color_AMD[0],\n)\nplt.plot(\n M_list, throughput_AMD_list, marker=\"o\", label=\"Real AMD MI210\", color=color_AMD[1]\n)\nplt.plot(\n M_list,\n throughput_AMD_sim_list,\n marker=\"x\",\n label=\"Simulated AMD MI210\",\n color=color_AMD[2],\n)\nplt.plot(\n M_list,\n throughput_TPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of Google TPUv3\",\n color=color_Google[0],\n)\nplt.plot(\n M_list,\n throughput_TPU_sim_list,\n marker=\"x\",\n label=\"Simulated Google TPUv3\",\n color=color_Google[2],\n)\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles, labels, loc=\"upper left\", bbox_to_anchor=(1, 1))\n# plt.legend()\n# plt.title(title)\nplt.xlabel(\"M\")\nplt.ylabel(\"G Elements/s\")\nplt.grid(True, which=\"both\", ls=\"--\", c=\"0.7\") # Adding a grid for better readability\nplt.savefig(\"figure5c.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)\n" }, { "path": "ae/figure5/cf/real_hardware/softmax_A100.csv", "content": "4096, 32, 9.99556025250909\n4096, 64, 19.634136210285714\n4096, 128, 37.27158060257627\n4096, 256, 68.719476736\n4096, 512, 111.34294964820253\n4096, 1024, 152.9755308210087\n4096, 2048, 167.5446289944381\n4096, 4096, 204.5603028420465\n4096, 8192, 216.51921285435077\n4096, 16384, 242.0249154863766\n4096, 32768, 240.88573103179803\n32, 4096, 8.457781752123077\n64, 4096, 16.65926708751515\n128, 4096, 30.97215852890141\n256, 4096, 57.11748715719481\n512, 4096, 96.6603628814066\n1024, 4096, 137.438953472\n2048, 4096, 166.75057862005687\n4096, 4096, 203.9673744280116\n8192, 4096, 229.96321626687583\n16384, 4096, 247.12465031664266\n32768, 4096, 278.826128490001" }, { "path": "ae/figure5/cf/real_hardware/softmax_MI210.csv", "content": "4096, 32, 5.389762881254902\n4096, 64, 10.372751205433962\n4096, 128, 13.49094021811043\n4096, 256, 36.049561566426235\n4096, 512, 60.662710497986204\n4096, 1024, 87.52331365381094\n4096, 2048, 75.18028224109402\n4096, 4096, 80.05545412703526\n4096, 8192, 70.97200623062432\n4096, 16384, 66.22940628486023\n4096, 32768, 61.59189862377593\n32, 4096, 4.822419420070175\n64, 4096, 9.012390391606559\n128, 4096, 15.820311191021583\n256, 4096, 28.28325730613505\n512, 4096, 43.76165682690547\n1024, 4096, 60.662710497986204\n2048, 4096, 72.54509709037526\n4096, 4096, 80.28379255865829\n8192, 4096, 84.78161949116145\n16384, 4096, 86.98237846435599\n32768, 4096, 88.01594018469544" }, { "path": "ae/figure5/cf/run.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../../..\n\npython -m ae.figure5.cf.test_softmax --simgpu --roofline\npython -m ae.figure5.cf.test_softmax --simtpu --roofline\npython -m ae.figure5.cf.test_softmax --simamd --roofline\n\npython -m ae.figure5.cf.test_softmax --simgpu\npython -m ae.figure5.cf.test_softmax --simtpu\npython -m ae.figure5.cf.test_softmax --simamd\n\ncd ae/figure5/cf\npython plot_softmax.py" }, { "path": "ae/figure5/cf/test_softmax.py", "content": "from software_model.softmax import Softmax\nfrom software_model.utils import data_type_dict, Tensor\nfrom hardware_model.device import device_dict\nimport argparse\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--gpu\", action=\"store_true\", help=\"Enable GPU\")\n parser.add_argument(\"--simgpu\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--simtpu\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--simamd\", action=\"store_true\", help=\"amd simulation\")\n parser.add_argument(\"--roofline\", action=\"store_true\", help=\"Roofline simulation\")\n args = parser.parse_args()\n\n A100 = device_dict[\"A100_80GB_fp16\"]\n TPU = device_dict[\"TPUv3\"]\n MI210 = device_dict[\"MI210\"]\n tpu_overhead = 300e-6\n gpu_overhead = 12e-6\n amd_overhead = MI210.compute_module.overhead.softmax\n\n if args.gpu:\n gpu_kernel_launch_overhead = Softmax.gpu_kernel_launch_overhead()\n\n print(f\"Performance of Softmax\")\n M = 2**12\n for N in range(5, 16):\n N = 2**N\n if args.simtpu:\n model = Softmax(data_type=data_type_dict[\"fp32\"])\n _ = model(Tensor([M, N], data_type=data_type_dict[\"fp32\"]))\n if args.roofline:\n latency = model.roofline_model(pcb_module=TPU) + tpu_overhead\n file_name = \"softmax_TPUv3_roofline.csv\"\n else:\n latency = model.compile_and_simulate(pcb_module=TPU) + tpu_overhead\n file_name = \"softmax_TPUv3_sim.csv\"\n else:\n model = Softmax(data_type=data_type_dict[\"fp16\"])\n _ = model(\n Tensor([M, N]),\n )\n if args.gpu:\n latency = model.run_on_gpu()\n if args.simgpu:\n if args.roofline:\n latency = model.roofline_model(pcb_module=A100) + gpu_overhead\n file_name = \"softmax_A100_roofline.csv\"\n else:\n latency = model.compile_and_simulate(pcb_module=A100) + gpu_overhead\n file_name = \"softmax_A100_sim.csv\"\n if args.simamd:\n model = Softmax(data_type=data_type_dict[\"fp32\"])\n _ = model(Tensor([M, N], data_type=data_type_dict[\"fp32\"]))\n if args.roofline:\n latency = model.roofline_model(pcb_module=MI210) + amd_overhead\n file_name = \"softmax_MI210_roofline.csv\"\n else:\n latency = (\n model.compile_and_simulate(pcb_module=MI210) + amd_overhead\n )\n file_name = \"softmax_MI210_sim.csv\"\n\n print(f\"{M}, {N}, {M*N/latency/1e9}\")\n with open(f\"ae/figure5/cf/{file_name}\", \"a\") as f:\n f.write(f\"{M}, {N}, {M*N/latency/1e9}\\n\")\n\n N = 2**12\n for M in range(5, 16):\n M = 2**M\n if args.simtpu:\n model = Softmax(data_type=data_type_dict[\"fp32\"])\n _ = model(Tensor([M, N], data_type=data_type_dict[\"fp32\"]))\n if args.roofline:\n latency = model.roofline_model(pcb_module=TPU) + tpu_overhead\n else:\n latency = model.compile_and_simulate(pcb_module=TPU) + tpu_overhead\n else:\n model = Softmax(data_type=data_type_dict[\"fp16\"])\n _ = model(\n Tensor([M, N]),\n )\n if args.gpu:\n latency = model.run_on_gpu()\n if args.simgpu:\n if args.roofline:\n latency = model.roofline_model(pcb_module=A100) + gpu_overhead\n else:\n latency = model.compile_and_simulate(pcb_module=A100) + gpu_overhead\n if args.simamd:\n model = Softmax(data_type=data_type_dict[\"fp32\"])\n _ = model(Tensor([M, N], data_type=data_type_dict[\"fp32\"]))\n if args.roofline:\n latency = model.roofline_model(pcb_module=MI210) + amd_overhead\n else:\n latency = (\n model.compile_and_simulate(pcb_module=MI210) + amd_overhead\n )\n print(f\"{M}, {N}, {M*N/latency/1e9}\")\n with open(f\"ae/figure5/cf/{file_name}\", \"a\") as f:\n f.write(f\"{M}, {N}, {M*N/latency/1e9}\\n\")\n" }, { "path": "ae/figure5/de/__init__.py", "content": "" }, { "path": "ae/figure5/de/plot_layernorm.py", "content": "import pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport numpy as np\n\n\nlayernorm_TPUv3_sim = pd.read_csv(\n \"layernorm_TPUv3_sim.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nlayernorm_TPUv3_sim.set_index([\"M\", \"N\"], inplace=True)\nlayernorm_TPUv3_roofline = pd.read_csv(\n \"layernorm_TPUv3_roofline.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nlayernorm_TPUv3_roofline.set_index([\"M\", \"N\"], inplace=True)\nlayernorm_A100 = pd.read_csv(\n \"real_hardware/layernorm_A100.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nlayernorm_A100.set_index([\"M\", \"N\"], inplace=True)\nlayernorm_A100_sim = pd.read_csv(\n \"layernorm_A100_sim.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nlayernorm_A100_sim.set_index([\"M\", \"N\"], inplace=True)\nlayernorm_A100_roofline = pd.read_csv(\n \"layernorm_A100_roofline.csv\", header=None, names=[\"M\", \"N\", \"throughput\"]\n)\nlayernorm_A100_roofline.set_index([\"M\", \"N\"], inplace=True)\n\n\ncolor_NV = sns.color_palette(\"Greens_d\", 4)[1:]\ncolor_Google = sns.color_palette(\"Blues_d\", 4)[1:]\ncolor_AMD = sns.color_palette(\"flare\", 3)\n\nM = 4096\ntitle = f\"Performance of layernorm with M={M}\"\nN_list = []\nthroughput_TPU_list = []\nthroughput_TPU_sim_list = []\nthroughput_TPU_roofline_list = []\nthroughput_GPU_list = []\nthroughput_GPU_sim_list = []\nthroughput_GPU_roofline_list = []\n\nfor N in range(6, 16):\n N = 2**N\n N_list.append(N)\n throughput_TPU_sim_list.append(\n layernorm_TPUv3_sim.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_TPU_roofline_list.append(\n layernorm_TPUv3_roofline.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_GPU_list.append(layernorm_A100.loc[(M, N), \"throughput\"].values[0])\n throughput_GPU_sim_list.append(\n layernorm_A100_sim.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_GPU_roofline_list.append(\n layernorm_A100_roofline.loc[(M, N), \"throughput\"].values[0]\n )\n\nplt.figure(figsize=(3.64, 2))\nplt.xscale(\"log\", base=2)\nplt.plot(\n N_list,\n throughput_GPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of NVIDIA A100\",\n color=color_NV[0],\n)\nplt.plot(\n N_list, throughput_GPU_list, marker=\"o\", label=\"Real NVIDIA A100\", color=color_NV[1]\n)\nplt.plot(\n N_list,\n throughput_GPU_sim_list,\n marker=\"x\",\n label=\"Simulated NVIDIA A100\",\n color=color_NV[2],\n)\nplt.plot(\n N_list,\n throughput_TPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of Google TPUv3\",\n color=color_Google[0],\n)\n\nplt.plot(\n N_list,\n throughput_TPU_sim_list,\n marker=\"x\",\n label=\"Simulated Google TPUv3\",\n color=color_Google[2],\n)\n\n\n# plt.legend()\n# plt.title(title)\nplt.xlabel(\"N\")\nplt.ylabel(\"G Elements/s\")\nplt.grid(True, which=\"both\", ls=\"--\", c=\"0.7\") # Adding a grid for better readability\nplt.savefig(\"figure5e.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)\n\nN = 4096\ntitle = f\"Performance of layernorm with N={N}\"\nM_list = []\nthroughput_TPU_list = []\nthroughput_TPU_sim_list = []\nthroughput_TPU_roofline_list = []\nthroughput_GPU_list = []\nthroughput_GPU_sim_list = []\nthroughput_GPU_roofline_list = []\n\nfor M in range(6, 16):\n M = 2**M\n M_list.append(M)\n throughput_TPU_sim_list.append(\n layernorm_TPUv3_sim.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_TPU_roofline_list.append(\n layernorm_TPUv3_roofline.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_GPU_list.append(layernorm_A100.loc[(M, N), \"throughput\"].values[0])\n throughput_GPU_sim_list.append(\n layernorm_A100_sim.loc[(M, N), \"throughput\"].values[0]\n )\n throughput_GPU_roofline_list.append(\n layernorm_A100_roofline.loc[(M, N), \"throughput\"].values[0]\n )\n\n\nplt.figure(figsize=(3.64, 2))\nplt.xscale(\"log\", base=2)\nplt.plot(\n M_list,\n throughput_GPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of NVIDIA A100\",\n color=color_NV[0],\n)\nplt.plot(\n M_list, throughput_GPU_list, marker=\"o\", label=\"Real NVIDIA A100\", color=color_NV[1]\n)\nplt.plot(\n M_list,\n throughput_GPU_sim_list,\n marker=\"x\",\n label=\"Simulated NVIDIA A100\",\n color=color_NV[2],\n)\nplt.plot(\n M_list,\n throughput_TPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of Google TPUv3\",\n color=color_Google[0],\n)\n\nplt.plot(\n M_list,\n throughput_TPU_sim_list,\n marker=\"x\",\n label=\"Simulated Google TPUv3\",\n color=color_Google[2],\n)\n\n# plt.legend()\n# plt.title(title)\nplt.xlabel(\"M\")\nplt.ylabel(\"G Elements/s\")\nplt.grid(True, which=\"both\", ls=\"--\", c=\"0.7\") # Adding a grid for better readability\nplt.savefig(\"figure5d.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)\n" }, { "path": "ae/figure5/de/real_hardware/layernorm_A100.csv", "content": "4096, 32, 2.476377540036036\n4096, 64, 4.8436635584845815\n4096, 128, 10.064179659276888\n4096, 256, 19.590407621844097\n4096, 512, 39.26827242057143\n4096, 1024, 78.888726656574\n4096, 2048, 139.06866438273516\n4096, 4096, 211.31755008307508\n4096, 8192, 265.2921552409576\n4096, 16384, 210.44858071824746\n4096, 32768, 199.2743197951547\n32, 4096, 2.4988900631272726\n64, 4096, 5.020601040073059\n128, 4096, 9.99556025250909\n256, 4096, 20.082404160292235\n512, 4096, 40.16480832058447\n1024, 4096, 79.6026517846878\n2048, 4096, 141.87246810012903\n4096, 4096, 211.6353208350797\n8192, 4096, 284.89370112414576\n16384, 4096, 345.7923546813956\n32768, 4096, 385.84643825998086" }, { "path": "ae/figure5/de/real_hardware/layernorm_MI210.csv", "content": "4096, 32, 2.3695247806698823\n4096, 64, 4.729122950631834\n4096, 128, 9.17590195596949\n4096, 256, 18.43588592802619\n4096, 512, 31.59641013907877\n4096, 1024, 52.44807583429779\n4096, 2048, 83.74459902278055\n4096, 4096, 89.41186119969811\n4096, 8192, 86.39576298097589\n4096, 16384, 86.3601901230308\n4096, 32768, 78.76990948928942\n32, 4096, 3.885694391427885\n64, 4096, 6.742243759280835\n128, 4096, 11.78431162745035\n256, 4096, 22.64512970663591\n512, 4096, 39.53264699736632\n1024, 4096, 58.88026281772146\n2048, 4096, 75.69152349271899\n4096, 4096, 89.32684940809148\n8192, 4096, 97.43158473170392\n16384, 4096, 101.13477801032278\n32768, 4096, 104.08326540288465" }, { "path": "ae/figure5/de/run.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../../..\n\npython -m ae.figure5.de.test_layernorm --simgpu --roofline\npython -m ae.figure5.de.test_layernorm --simtpu --roofline\n\npython -m ae.figure5.de.test_layernorm --simgpu\npython -m ae.figure5.de.test_layernorm --simtpu\n\ncd ae/figure5/de\npython plot_layernorm.py" }, { "path": "ae/figure5/de/test_layernorm.py", "content": "from software_model.layernorm import LayerNorm\nfrom software_model.utils import data_type_dict, Tensor\nfrom hardware_model.device import device_dict\nimport argparse\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--gpu\", action=\"store_true\", help=\"Enable GPU\")\n parser.add_argument(\"--amd\", action=\"store_true\", help=\"Enable AMD\")\n parser.add_argument(\"--simgpu\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--simtpu\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--simamd\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--roofline\", action=\"store_true\", help=\"Roofline simulation\")\n args = parser.parse_args()\n\n A100 = device_dict[\"A100_80GB_fp16\"]\n TPU = device_dict[\"TPUv3\"]\n MI210 = device_dict[\"MI210\"]\n\n if args.gpu:\n gpu_kernel_launch_overhead = LayerNorm.gpu_kernel_launch_overhead()\n\n print(f\"Performance of LayerNorm\")\n M = 2**12\n for N in range(5, 16):\n N = 2**N\n # N = 2**15\n if args.simtpu:\n model = LayerNorm(data_type=data_type_dict[\"fp32\"])\n _ = model(Tensor([M, N], data_type=data_type_dict[\"fp32\"]))\n if args.roofline:\n latency = model.roofline_model(TPU) + 140e-6\n file_name = \"layernorm_TPUv3_roofline.csv\"\n else:\n latency = (\n model.compile_and_simulate(\n pcb_module=TPU, compile_mode=\"heuristic-TPU\"\n )\n + 140e-6\n )\n file_name = \"layernorm_TPUv3_sim.csv\"\n else:\n model = LayerNorm(data_type=data_type_dict[\"fp16\"])\n _ = model(\n Tensor([M, N]),\n )\n if args.gpu:\n latency = model.run_on_gpu()\n if args.amd:\n # model.amd_kernel_launch_overhead()\n latency = model.run_on_amd()\n if args.simgpu:\n if args.roofline:\n latency = model.roofline_model(A100) + 4.5e-5\n file_name = \"layernorm_A100_roofline.csv\"\n else:\n latency = (\n model.compile_and_simulate(\n pcb_module=A100, compile_mode=\"heuristic-GPU\"\n )\n + 4.5e-5\n )\n file_name = \"layernorm_A100_sim.csv\"\n if args.simamd:\n model = LayerNorm(data_type=data_type_dict[\"fp32\"])\n _ = model(Tensor([M, N], data_type=data_type_dict[\"fp32\"]))\n if args.roofline:\n latency = (\n model.roofline_model(MI210)\n + MI210.compute_module.overhead.layernorm\n )\n file_name = \"layernorm_MI210_roofline.csv\"\n else:\n latency = (\n model.compile_and_simulate(\n pcb_module=MI210, compile_mode=\"heuristic-GPU\"\n )\n + MI210.compute_module.overhead.layernorm\n )\n file_name = \"layernorm_MI210_sim.csv\"\n print(f\"{M}, {N}, {M*N/latency/1e9}\")\n with open(f\"ae/figure5/de/{file_name}\", \"a\") as f:\n f.write(f\"{M}, {N}, {M*N/latency/1e9}\\n\")\n\n N = 2**12\n for M in range(5, 16):\n M = 2**M\n if args.simtpu:\n model = LayerNorm(data_type=data_type_dict[\"fp32\"])\n _ = model(Tensor([M, N], data_type=data_type_dict[\"fp32\"]))\n if args.roofline:\n latency = model.roofline_model(TPU) + 140e-6\n else:\n latency = (\n model.compile_and_simulate(\n pcb_module=TPU, compile_mode=\"heuristic-TPU\"\n )\n + 140e-6\n )\n else:\n model = LayerNorm(data_type=data_type_dict[\"fp16\"])\n _ = model(\n Tensor([M, N]),\n )\n if args.gpu:\n latency = model.run_on_gpu()\n if args.amd:\n latency = model.run_on_amd()\n if args.simgpu:\n if args.roofline:\n latency = model.roofline_model(A100) + 4.5e-5\n else:\n latency = (\n model.compile_and_simulate(\n pcb_module=A100, compile_mode=\"heuristic-GPU\"\n )\n + 4.5e-5\n )\n if args.simamd:\n model = LayerNorm(data_type=data_type_dict[\"fp32\"])\n _ = model(Tensor([M, N], data_type=data_type_dict[\"fp32\"]))\n if args.roofline:\n latency = (\n model.roofline_model(MI210)\n + MI210.compute_module.overhead.layernorm\n )\n else:\n latency = (\n model.compile_and_simulate(\n pcb_module=MI210, compile_mode=\"heuristic-GPU\"\n )\n + MI210.compute_module.overhead.layernorm\n )\n print(f\"{M}, {N}, {M*N/latency/1e9}\")\n with open(f\"ae/figure5/de/{file_name}\", \"a\") as f:\n f.write(f\"{M}, {N}, {M*N/latency/1e9}\\n\")\n" }, { "path": "ae/figure5/g/__init__.py", "content": "" }, { "path": "ae/figure5/g/plot_gelu.py", "content": "import pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport numpy as np\n\n\ngelu_TPUv3_sim = pd.read_csv(\n \"gelu_TPUv3_sim.csv\", header=None, names=[\"M\", \"throughput\"]\n)\ngelu_TPUv3_roofline = pd.read_csv(\n \"gelu_TPUv3_roofline.csv\", header=None, names=[\"M\", \"throughput\"]\n)\ngelu_A100 = pd.read_csv(\"real_hardware/gelu_A100.csv\", header=None, names=[\"M\", \"throughput\"])\ngelu_A100_sim = pd.read_csv(\"gelu_A100_sim.csv\", header=None, names=[\"M\", \"throughput\"])\ngelu_A100_roofline = pd.read_csv(\n \"gelu_A100_roofline.csv\", header=None, names=[\"M\", \"throughput\"]\n)\ngelu_MI210 = pd.read_csv(\"real_hardware/gelu_MI210.csv\", header=None, names=[\"M\", \"throughput\"])\ngelu_MI210_sim = pd.read_csv(\n \"gelu_MI210_sim.csv\", header=None, names=[\"M\", \"throughput\"]\n)\ngelu_MI210_roofline = pd.read_csv(\n \"gelu_MI210_roofline.csv\", header=None, names=[\"M\", \"throughput\"]\n)\n\ncolor_NV = sns.color_palette(\"Greens_d\", 4)[1:]\ncolor_Google = sns.color_palette(\"Blues_d\", 4)[1:]\ncolor_AMD = sns.color_palette(\"flare\", 3)\n\nM = 4096\ntitle = f\"Performance of gelu with M={M}\"\nM_list = []\nthroughput_TPU_list = []\nthroughput_TPU_sim_list = []\nthroughput_TPU_roofline_list = []\nthroughput_GPU_list = []\nthroughput_GPU_sim_list = []\nthroughput_GPU_roofline_list = []\nthroughput_AMD_list = []\nthroughput_AMD_sim_list = []\nthroughput_AMD_roofline_list = []\nfor M in range(10, 30):\n M = 2**M\n M_list.append(M)\n \n throughput_TPU_sim_list.append(\n gelu_TPUv3_sim[gelu_TPUv3_sim[\"M\"] == M][\"throughput\"].iloc[0]\n )\n throughput_TPU_roofline_list.append(\n gelu_TPUv3_roofline[gelu_TPUv3_roofline[\"M\"] == M][\"throughput\"].iloc[0]\n )\n throughput_GPU_list.append(gelu_A100[gelu_A100[\"M\"] == M][\"throughput\"].iloc[0])\n throughput_GPU_sim_list.append(\n gelu_A100_sim[gelu_A100_sim[\"M\"] == M][\"throughput\"].iloc[0]\n )\n throughput_GPU_roofline_list.append(\n gelu_A100_roofline[gelu_A100_roofline[\"M\"] == M][\"throughput\"].iloc[0]\n )\n throughput_AMD_list.append(gelu_MI210[gelu_MI210[\"M\"] == M][\"throughput\"].iloc[0])\n throughput_AMD_sim_list.append(\n gelu_MI210_sim[gelu_MI210_sim[\"M\"] == M][\"throughput\"].iloc[0]\n )\n throughput_AMD_roofline_list.append(\n gelu_MI210_roofline[gelu_MI210_roofline[\"M\"] == M][\"throughput\"].iloc[0]\n )\n\nplt.figure(figsize=(6, 2.3))\nplt.xscale(\"log\", base=2)\nplt.plot(\n M_list,\n throughput_GPU_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of NVIDIA A100\",\n color=color_NV[0],\n)\nplt.plot(\n M_list, throughput_GPU_list, marker=\"o\", label=\"Real NVIDIA A100\", color=color_NV[1]\n)\nplt.plot(\n M_list,\n throughput_GPU_sim_list,\n marker=\"x\",\n label=\"Simulated NVIDIA A100\",\n color=color_NV[2],\n)\nplt.plot(\n M_list,\n throughput_AMD_roofline_list,\n marker=\" \",\n linewidth=1.5,\n linestyle=\"--\",\n label=\"Roofline of AMD MI210\",\n color=color_AMD[0],\n)\nplt.plot(\n M_list, throughput_AMD_list, marker=\"o\", label=\"Real AMD MI210\", color=color_AMD[1]\n)\nplt.plot(\n M_list,\n throughput_AMD_sim_list,\n marker=\"x\",\n label=\"Simulated AMD MI210\",\n color=color_AMD[2],\n)\nplt.plot(\n M_list,\n throughput_TPU_roofline_list,\n marker=\" \",\n linewidth=3.5,\n linestyle=\"--\",\n label=\"Roofline of Google TPUv3\",\n color=color_Google[0],\n)\nplt.plot(\n M_list,\n throughput_TPU_sim_list,\n marker=\"x\",\n label=\"Simulated Google TPUv3\",\n color=color_Google[2],\n)\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles, labels, loc=\"upper left\", bbox_to_anchor=(1, 1.05))\n# plt.title(title)\nplt.xlabel(\"# Elements\")\nplt.ylabel(\"G Elements/s\")\nplt.grid(True, which=\"both\", ls=\"--\", c=\"0.7\") # Adding a grid for better readability\nplt.savefig(f\"figure5g.pdf\", dpi=300, bbox_inches=\"tight\", pad_inches=0.01)\n" }, { "path": "ae/figure5/g/real_hardware/gelu_A100.csv", "content": "1024, 0.021262214336633663\n2048, 0.04338350804040404\n4096, 0.08547198599004975\n8192, 0.1700977146930693\n16384, 0.3368601800784314\n32768, 0.6889170600100251\n65536, 1.38129601479397\n131072, 2.735103551681592\n262144, 5.52518405917588\n524288, 11.19095804352163\n1048576, 21.935394070344138\n2097152, 43.980465111040004\n4194304, 86.66101499712316\n8388608, 141.87246810012903\n16777216, 210.36993775086398\n33554432, 277.042299912063\n67108864, 328.2507017033889\n134217728, 360.17271492086496\n268435456, 379.53814489891255\n536870912, 390.32758080867535" }, { "path": "ae/figure5/g/real_hardware/gelu_MI210.csv", "content": "1024, 0.047197442813186816\n2048, 0.09761289309090909\n4096, 0.19522578618181818\n8192, 0.39045157236363637\n16384, 0.7898790429425288\n32768, 1.5618062894545455\n65536, 3.0885158083595505\n131072, 6.108397932088889\n262144, 11.822705675010754\n524288, 21.55905152501961\n1048576, 37.27158060257627\n2097152, 59.43306096086487\n4194304, 87.96093022208001\n8388608, 114.60707520792182\n16777216, 133.52702880012146\n33554432, 146.984322042118\n67108864, 151.08694402074934\n134217728, 152.91320207016489\n268435456, 154.04294798777178\n536870912, 154.44972829556897" }, { "path": "ae/figure5/g/run.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../../..\n\npython -m ae.figure5.g.test_gelu --simgpu --roofline\npython -m ae.figure5.g.test_gelu --simtpu --roofline\npython -m ae.figure5.g.test_gelu --simamd --roofline\n\npython -m ae.figure5.g.test_gelu --simgpu\npython -m ae.figure5.g.test_gelu --simtpu\npython -m ae.figure5.g.test_gelu --simamd\n\ncd ae/figure5/g\npython plot_gelu.py" }, { "path": "ae/figure5/g/test_gelu.py", "content": "from software_model.gelu import GeLU\nfrom software_model.utils import data_type_dict, Tensor\nfrom hardware_model.device import device_dict\nimport argparse\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--gpu\", action=\"store_true\", help=\"Enable GPU\")\n parser.add_argument(\"--simgpu\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--simtpu\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--roofline\", action=\"store_true\", help=\"Roofline simulation\")\n parser.add_argument(\"--amd\", action=\"store_true\", help=\"Enable AMD\")\n parser.add_argument(\"--simamd\", action=\"store_true\", help=\"Enable simulation\")\n args = parser.parse_args()\n\n A100 = device_dict[\"A100_80GB_fp16\"]\n TPU = device_dict[\"TPUv3\"]\n MI210 = device_dict[\"MI210\"]\n\n if args.gpu:\n gpu_kernel_launch_overhead = GeLU.gpu_kernel_launch_overhead()\n\n print(f\"Performance of GELU\")\n for M in range(10, 30):\n M = 2**M\n # N = 2**15\n if args.simtpu:\n model = GeLU(data_type=data_type_dict[\"fp32\"])\n _ = model(Tensor([M], data_type=data_type_dict[\"fp32\"]))\n if args.roofline:\n latency = model.roofline_model(TPU) + 100e-6\n file_name = \"gelu_TPUv3_roofline.csv\"\n else:\n latency = (\n model.compile_and_simulate(\n pcb_module=TPU, compile_mode=\"heuristic-TPU\"\n )\n + 100e-6\n )\n file_name = \"gelu_TPUv3_sim.csv\"\n else:\n model = GeLU(data_type=data_type_dict[\"fp16\"])\n _ = model(\n Tensor([M]),\n )\n if args.gpu:\n latency = model.run_on_gpu()\n if args.amd:\n model.amd_kernel_launch_overhead()\n latency = model.run_on_amd()\n if args.simgpu:\n if args.roofline:\n latency = model.roofline_model(A100) + 4.5e-5\n file_name = \"gelu_A100_roofline.csv\"\n else:\n latency = (\n model.compile_and_simulate(\n pcb_module=A100, compile_mode=\"heuristic-GPU\"\n )\n + 4.5e-5\n )\n file_name = \"gelu_A100_sim.csv\"\n if args.simamd:\n if args.roofline:\n latency = (\n model.roofline_model(MI210) + MI210.compute_module.overhead.gelu\n )\n file_name = \"gelu_MI210_roofline.csv\"\n else:\n latency = (\n model.compile_and_simulate(\n pcb_module=MI210, compile_mode=\"heuristic-GPU\"\n )\n + MI210.compute_module.overhead.gelu\n )\n file_name = \"gelu_MI210_sim.csv\"\n print(f\"{M}, {M/latency/1e9}\")\n with open(f\"ae/figure5/g/{file_name}\", \"a\") as f:\n f.write(f\"{M}, {M/latency/1e9}\\n\")\n" }, { "path": "ae/figure5/h/__init__.py", "content": "" }, { "path": "ae/figure5/h/run.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../../..\n\npython -m ae.figure5.h.test_allreduce\n" }, { "path": "ae/figure5/h/test_allreduce.py", "content": "from software_model.communication_primitives import AllReduceMultiPCB\nfrom software_model.utils import data_type_dict, Tensor\nfrom hardware_model.interconnect import interconnect_module_dict\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport numpy as np\n\nif __name__ == \"__main__\":\n interconnect_module = interconnect_module_dict[\"NVLinkV3_FC_4\"]\n gpu_latency_list = [\n 12.52,\n 13.92,\n 12.39,\n 13.22,\n 12.35,\n 12.45,\n 13.12,\n 13.02,\n 15.12,\n 15.23,\n 15.99,\n 17.39,\n 20.00,\n 22.93,\n 28.66,\n 35.93,\n 47.27,\n 60.75,\n 66.40,\n 84.75,\n 128.8,\n 195.7,\n 279.7,\n 532.3,\n 961.7,\n 1883.7,\n 3659.0,\n 7219.2,\n 14136,\n 27944,\n 55384,\n 110277,\n ]\n size_list = [\n 8,\n 16,\n 32,\n 64,\n 128,\n 256,\n 512,\n 1024,\n 2048,\n 4096,\n 8192,\n 16384,\n 32768,\n 65536,\n 131072,\n 262144,\n 524288,\n 1048576,\n 2097152,\n 4194304,\n 8388608,\n 16777216,\n 33554432,\n 67108864,\n 134217728,\n 268435456,\n 536870912,\n 1073741824,\n 2147483648,\n 4294967296,\n 8589934592,\n 17179869184,\n ]\n simulated_latency_list = []\n data_type = data_type_dict[\"fp16\"]\n for data_size in size_list:\n model = AllReduceMultiPCB(data_type=data_type)\n _ = model(\n Tensor([data_size / 2]),\n )\n our_latency = model.simulate(interconnect_module=interconnect_module)\n simulated_latency_list.append(our_latency * 1e6)\n\n gpu_bandwidth_list = np.array(size_list) / np.array(gpu_latency_list) / 1e3\n simulated_gpu_bandwidth_list = (\n np.array(size_list) / np.array(simulated_latency_list) / 1e3\n )\n\n size_list = size_list[9:]\n gpu_bandwidth_list = gpu_bandwidth_list[9:]\n simulated_gpu_bandwidth_list = simulated_gpu_bandwidth_list[9:]\n\n color_NV = sns.color_palette(\"Greens_d\", 4)[1:]\n color_Google = sns.color_palette(\"Blues_d\", 4)[1:]\n\n plt.figure(figsize=(6, 2.3))\n plt.xscale(\"log\", base=2)\n plt.plot(\n size_list,\n gpu_bandwidth_list,\n marker=\"o\",\n label=\"Real NVIDIA A100 Node\",\n color=color_NV[0],\n )\n plt.plot(\n size_list,\n simulated_gpu_bandwidth_list,\n marker=\"x\",\n label=\"Simulated NVIDIA A100 Node\",\n color=color_NV[2],\n )\n\n interconnect_module = interconnect_module_dict[\"TPUv3Link_8\"]\n simulated_tpu_bandwidth_list = []\n data_type = data_type_dict[\"fp16\"]\n for data_size in size_list:\n model = AllReduceMultiPCB(data_type=data_type)\n _ = model(\n Tensor([data_size // 2]),\n )\n our_latency = model.simulate(interconnect_module=interconnect_module)\n simulated_tpu_bandwidth_list.append(data_size / our_latency / 1e9)\n\n \n plt.plot(\n size_list,\n simulated_tpu_bandwidth_list,\n marker=\"x\",\n label=\"Simulated Google TPU v3 Node\",\n color=color_Google[2],\n )\n plt.xlabel(\"Data Size (Bytes)\")\n plt.ylabel(\"Bandwidth (GB/s)\")\n plt.grid(\n True, which=\"both\", ls=\"--\", c=\"0.7\"\n ) # Adding a grid for better readability\n plt.legend()\n plt.savefig(\n \"ae/figure5/h/figure5h.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300\n )\n" }, { "path": "ae/figure5/ijkl/__init__.py", "content": "" }, { "path": "ae/figure5/ijkl/plot_transformer.py", "content": "import matplotlib.pyplot as plt\nimport seaborn as sns\nimport csv\nimport pandas as pd\n\ndef read_csv(filename: str):\n numbers = []\n # Open the CSV file and read the numbers\n with open(filename, \"r\") as csvfile:\n reader = csv.reader(csvfile)\n for row in reader:\n # Since each row contains only one number, we use row[0]\n numbers.append(float(row[0]))\n return numbers\n\n\ncategories = [\n \"Q_K_V\",\n \"Q_mul_K\",\n \"A_mul_V\",\n \"Wo_proj\",\n \"W1_proj\",\n \"W2_proj\",\n \"Softmax\",\n \"LayerNorm_MHA\",\n \"LayerNorm_FFN\",\n \"GeLU\",\n \"AllReduce_MHA\",\n \"AllReduce_FFN\",\n]\ncolors_matmul = sns.color_palette(\"flare_r\", 6)\ncolors_normalization = sns.color_palette(\"summer\", 3)\ncolors_gelu = sns.color_palette(\"pink\", 1)\ncolors_allreduce = sns.color_palette(\"Blues_r\", 2)\ncolors = colors_matmul + colors_normalization + colors_gelu + colors_allreduce\n# values_simgpu = read_csv(\"transformer_A100_sim.csv\")\nvalues_simgpu = pd.read_csv(\"transformer_A100_sim.csv\", header=None, names=categories, index_col=None).iloc[0].tolist()\nprint(values_simgpu)\nvalues_gpu = read_csv(\"real_hardware/transformer_A100.csv\")\n# values_gpu_roofline = read_csv(\"transformer_A100_roofline.csv\")\nvalues_gpu_roofline = pd.read_csv(\"transformer_A100_roofline.csv\", header=None, names=categories, index_col=None).iloc[0].tolist()\n# values_simtpu = read_csv(\"transformer_TPUv3_sim.csv\")\nvalues_simtpu = pd.read_csv(\"transformer_TPUv3_sim.csv\", header=None, names=categories, index_col=None).iloc[0].tolist()\n\n# values_tpu_roofline = read_csv(\"transformer_TPUv3_roofline.csv\")\nvalues_tpu_roofline = pd.read_csv(\"transformer_TPUv3_roofline.csv\", header=None, names=categories, index_col=None).iloc[0].tolist()\n\nplt.figure(figsize=(3, 2.8))\n\n# Create the stacked bar graph\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_gpu)):\n plt.bar(1, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n bottom += value\nvalue_gt = bottom\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_simgpu)):\n plt.bar(2, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\nvalue_sim = bottom\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_gpu_roofline)):\n plt.bar(3, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\nvalue_roofline = bottom\n\nprint(f\"gpu prefilling: {value_sim/value_gt}, {value_roofline/value_gt}\")\n# Set the title, legend, and display the graph\n# plt.title(\n# \"GPU Runtime Breakdown of One Transformer Layer in GPT-3 \\n(Initial computation, batch size = 8, sequence length = 2048)\"\n# )\nplt.ylabel(\"Latency (s)\")\n# plt.xlabel('Bar Sets')\nplt.xticks([1, 2, 3], [\"Real\\nA100\", \"Simulated\\nA100\", \"Roofline\\nModel\"])\n# handles, labels = plt.gca().get_legend_handles_labels()\n# plt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1, 1))\nplt.tight_layout()\nplt.savefig(\"figure5i.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)\n\n\nplt.figure(figsize=(3, 2.8))\n\n# Create the stacked bar graph\n# bottom = 0\n# for i, (category, value) in enumerate(zip(categories, values_tpu)):\n# plt.bar(1, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n# bottom += value\n# value_gt = bottom\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_simtpu)):\n plt.bar(2, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\nvalue_sim = bottom\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_tpu_roofline)):\n plt.bar(3, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\nvalue_roofline = bottom\n# print(f\"tpu prefilling: {value_sim/value_gt}, {value_roofline/value_gt}\")\n\n# Set the title, legend, and display the graph\n# plt.title(\n# \"TPU Runtime Breakdown of One Transformer Layer in GPT-3 \\n(Initial computation, batch size = 8, sequence length = 2048)\"\n# )\nplt.ylabel(\"Latency (s)\")\n# plt.xlabel('Bar Sets')\nplt.xticks([1, 2, 3], [\"Real\\nTPUv3\", \"Simulated\\nTPUv3\", \"Roofline\\nModel\"])\n# handles, labels = plt.gca().get_legend_handles_labels()\n# plt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1, 1))\nplt.tight_layout()\nplt.savefig(\"figure5j.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)\n\n\n# values_simgpu = read_csv(\"transformerAR_A100_sim.csv\")\nvalues_simgpu=pd.read_csv(\"transformerAR_A100_sim.csv\",header=None,names=categories,index_col=None).iloc[0].tolist()\nvalues_gpu = read_csv(\"real_hardware/transformerAR_A100.csv\")\n# values_gpu_roofline = read_csv(\"transformerAR_A100_roofline.csv\")\nvalues_gpu_roofline=pd.read_csv(\"transformerAR_A100_roofline.csv\",header=None,names=categories,index_col=None).iloc[0].tolist()\n\nplt.figure(figsize=(3, 2.8))\n\n# Create the stacked bar graph\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_gpu)):\n value = value * 1e3\n plt.bar(1, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n bottom += value\nvalue_gt = bottom\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_simgpu)):\n value = value * 1e3\n plt.bar(2, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\nvalue_sim = bottom\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_gpu_roofline)):\n value = value * 1e3\n plt.bar(3, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\nvalue_roofline = bottom\n\nprint(value_sim / value_gt, value_roofline / value_gt)\n\n# Set the title, legend, and display the graph\n# plt.title(\n# \"GPU Runtime Breakdown of One Transformer Layer in GPT-3 \\n(Auto regression, batch size = 8, sequence length = 2048)\"\n# )\nplt.ylabel(\"Latency (ms)\")\n# plt.xlabel('Bar Sets')\nplt.xticks([1, 2, 3], [\"Real\\nA100\", \"Simulated\\nA100\", \"Roofline\\nModel\"])\n# handles, labels = plt.gca().get_legend_handles_labels()\n# plt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1, 1))\nplt.tight_layout()\nplt.savefig(\"figure5k.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)\n\n\n\n\n# values_simtpu = read_csv(\"transformerAR_TPUv3_sim.csv\")\nvalues_simtpu = pd.read_csv(\"transformerAR_TPUv3_sim.csv\", header=None, names=categories, index_col=None).iloc[0].tolist()\n# values_tpu = read_csv(\"real_hardware/transformerAR_TPUv3.csv\")\n# values_tpu_roofline=read_csv(\"transformerAR_TPUv3_roofline.csv\")\nvalues_tpu_roofline=pd.read_csv(\"transformerAR_TPUv3_roofline.csv\",header=None,names=categories,index_col=None).iloc[0].tolist()\n\nplt.figure(figsize=(4.5, 2.8))\n\n# Create the stacked bar graph\n# bottom = 0\n# for i, (category, value) in enumerate(zip(categories, values_tpu)):\n# value=value*1e3\n# plt.bar(1, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n# bottom += value\n# value_gt=bottom\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_simtpu)):\n value=value*1e3\n plt.bar(2, value, bottom=bottom, color=colors[i], label=category,width=0.5)\n bottom += value\nvalue_sim=bottom\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_tpu_roofline)):\n value=value*1e3\n plt.bar(3, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\nvalue_roofline=bottom\nprint(value_sim/value_gt,value_roofline/value_gt)\n\n# Set the title, legend, and display the graph\n# plt.title(\n# \"GPU Runtime Breakdown of One Transformer Layer in GPT-3 \\n(Auto regression, batch size = 8, input(output) sequence length = 2048(1024))\"\n# )\nplt.ylabel(\"Latency (ms)\")\n# plt.xlabel('Bar Sets')\nplt.xticks([1, 2, 3], [\"Real\\nTPUv3\", \"Simulated\\nTPUv3\", \"Roofline\\nModel\"])\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1, 1.05))\nplt.tight_layout()\nplt.savefig(\"figure5l.pdf\",bbox_inches=\"tight\", pad_inches=0.01, dpi=300)" }, { "path": "ae/figure5/ijkl/real_hardware/transformerAR_A100.csv", "content": "0.0002124309539794922\n0.00010609626770019531\n0.0001386404037475586\n6.890296936035156e-05\n0.00018596649169921875\n0.00018608570098876953\n1.6689300537109375e-05\n4.8041343688964844e-05\n4.8041343688964844e-05\n4.792213439941406e-05\n26.04e-06\n26.04e-06" }, { "path": "ae/figure5/ijkl/real_hardware/transformer_A100.csv", "content": "0.013721823692321777\n0.0018811225891113281\n0.001183152198791504\n0.0045403242111206055\n0.017464280128479004\n0.017485618591308594\n0.00280153751373291\n0.0006816387176513672\n0.0006816387176513672\n0.0005242824554443359\n0.0028909\n0.0028909" }, { "path": "ae/figure5/ijkl/run.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../../..\n\npython -m ae.figure5.ijkl.test_transformer --simgpu --roofline\npython -m ae.figure5.ijkl.test_transformer --simtpu --roofline\npython -m ae.figure5.ijkl.test_transformer --simgpu --init --roofline\npython -m ae.figure5.ijkl.test_transformer --simtpu --init --roofline\n\npython -m ae.figure5.ijkl.test_transformer --simgpu\npython -m ae.figure5.ijkl.test_transformer --simtpu\npython -m ae.figure5.ijkl.test_transformer --simgpu --init\npython -m ae.figure5.ijkl.test_transformer --simtpu --init\n\ncd ae/figure5/ijkl\npython plot_transformer.py" }, { "path": "ae/figure5/ijkl/test_transformer.py", "content": "from software_model.transformer import (\n TransformerBlockInitComputationTP,\n TransformerBlockAutoRegressionTP,\n)\nfrom software_model.utils import data_type_dict, Tensor\nfrom hardware_model.system import system_dict\nimport argparse\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--init\", action=\"store_true\", help=\"initial computation\")\n parser.add_argument(\"--gpu\", action=\"store_true\", help=\"Enable GPU\")\n parser.add_argument(\"--simgpu\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--simtpu\", action=\"store_true\", help=\"Enable simulation\")\n parser.add_argument(\"--roofline\", action=\"store_true\", help=\"use roofline\")\n args = parser.parse_args()\n\n bs = 8\n s = 2048\n if args.init:\n print(\"Initial computation\")\n if args.simgpu:\n model = TransformerBlockInitComputationTP(\n d_model=12288,\n n_heads=96,\n device_count=4,\n data_type=data_type_dict[\"fp16\"],\n )\n A100_system = system_dict[\"A100_4_fp16\"]\n # from design_space_exploration.dse import read_architecture_template, template_to_system\n # arch_specs = read_architecture_template(\"configs/template.json\")\n # A100_system = template_to_system(arch_specs)\n _ = model(Tensor([bs, s, 12288], data_type_dict[\"fp16\"]))\n if args.roofline:\n model.roofline_model(A100_system)\n file_name = \"transformer_A100_roofline.csv\"\n else:\n model.compile_and_simulate(A100_system, compile_mode=\"heuristic-GPU\")\n file_name = \"transformer_A100_sim.csv\"\n if args.simtpu:\n model = TransformerBlockInitComputationTP(\n d_model=12288,\n n_heads=96,\n device_count=8,\n data_type=data_type_dict[\"fp16\"],\n )\n TPU_system = system_dict[\"TPUv3_8\"]\n _ = model(Tensor([bs, s, 12288], data_type_dict[\"fp16\"]))\n if args.roofline:\n model.roofline_model(TPU_system)\n file_name = \"transformer_TPUv3_roofline.csv\"\n else:\n model.compile_and_simulate(TPU_system, compile_mode=\"heuristic-TPU\")\n file_name = \"transformer_TPUv3_sim.csv\"\n if args.gpu:\n model = TransformerBlockInitComputationTP(\n d_model=12288,\n n_heads=96,\n device_count=4,\n data_type=data_type_dict[\"fp16\"],\n )\n _ = model(Tensor([bs, s, 12288], data_type_dict[\"fp16\"]))\n model.run_on_gpu()\n else:\n print(\"Auto-regression\")\n output_token_length = 1024\n if args.simgpu:\n model = TransformerBlockAutoRegressionTP(\n d_model=12288,\n n_heads=96,\n device_count=4,\n data_type=data_type_dict[\"fp16\"],\n )\n A100_system = system_dict[\"A100_4_fp16\"]\n _ = model(\n Tensor([bs, 1, 12288], data_type_dict[\"fp16\"]), s + output_token_length\n )\n if args.roofline:\n model.roofline_model(A100_system)\n file_name = \"transformerAR_A100_roofline.csv\"\n else:\n model.compile_and_simulate(A100_system, compile_mode=\"heuristic-GPU\")\n file_name = \"transformerAR_A100_sim.csv\"\n if args.simtpu:\n model = TransformerBlockAutoRegressionTP(\n d_model=12288,\n n_heads=96,\n device_count=8,\n data_type=data_type_dict[\"fp16\"],\n )\n TPU_system = system_dict[\"TPUv3_8\"]\n _ = model(\n Tensor([bs, 1, 12288], data_type_dict[\"fp16\"]), s + output_token_length\n )\n if args.roofline:\n model.roofline_model(TPU_system)\n file_name = \"transformerAR_TPUv3_roofline.csv\"\n else:\n model.compile_and_simulate(TPU_system, compile_mode=\"heuristic-TPU\")\n file_name = \"transformerAR_TPUv3_sim.csv\"\n if args.gpu:\n model = TransformerBlockAutoRegressionTP(\n d_model=12288,\n n_heads=96,\n device_count=4,\n data_type=data_type_dict[\"fp16\"],\n )\n _ = model(\n Tensor([bs, 1, 12288], data_type_dict[\"fp16\"]), s + output_token_length\n )\n model.run_on_gpu()\n with open(f\"ae/figure5/ijkl/{file_name}\", \"w\") as f:\n if args.roofline:\n f.write(model.roofline_log)\n else:\n f.write(model.simluate_log)\n" }, { "path": "ae/figure5/run_figure5.sh", "content": "cd ab\nbash run.sh\ncd ..\n\ncd cf\nbash run.sh\ncd ..\n\ncd de\nbash run.sh\ncd ..\n\ncd g\nbash run.sh\ncd ..\n\ncd h\nbash run.sh\ncd ..\n\ncd ijkl\nbash run.sh\ncd .." }, { "path": "ae/figure6/real_hardware/die_area.csv", "content": "476.25,\t446.22\n76.44,\t33\n119.31,\t25.2\n58,\t83.26\n31.77,\t40.83\n20.95,\t45.52\n0,\t42\n40,\t4" }, { "path": "ae/figure6/run_figure6.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../..\n\npython -m ae.figure6.test_cost_model" }, { "path": "ae/figure6/test_cost_model.py", "content": "from cost_model.cost_model import calc_compute_chiplet_area_mm2, calc_io_die_area_mm2\nfrom design_space_exploration.dse import read_architecture_template\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport pandas as pd\n\narch_specs = read_architecture_template(\"configs/GA100.json\")\n\ncompute_chiplet_area_mm2, A100_core_breakdown_map, compute_total_die_map = (\n calc_compute_chiplet_area_mm2(arch_specs, verbose=True)\n)\nio_die_area_mm2, io_total_die_map = calc_io_die_area_mm2(arch_specs, verbose=True)\n\n\ncategories = [\n \"Cores\",\n \"On-chip interconnect\",\n \"Global buffer\",\n \"Memory(PHY)\",\n \"Memory(Control)\",\n \"Off-chip interconnect\\n(PHY)\",\n \"Off-chip interconnect\\n(Control)\",\n \"Other\",\n]\n\ndie_area = pd.read_csv(\n \"ae/figure6/real_hardware/die_area.csv\", header=None, names=[\"A100\", \"MI210\"]\n)\n\nvalues_a100 = die_area[\"A100\"].tolist()\nvalues_mi210 = die_area[\"MI210\"].tolist()\n\nvalues_a100_sim = [\n compute_total_die_map[\"cores_area\"],\n compute_total_die_map[\"crossbar_area\"],\n io_total_die_map[\"global_buffer_area\"],\n io_total_die_map[\"mem_phy_area\"],\n io_total_die_map[\"mem_controller_area\"],\n io_total_die_map[\"device_phy_area\"],\n io_total_die_map[\"device_controller_area\"],\n 0,\n]\n\n\narch_specs = read_architecture_template(\"configs/mi210_template.json\")\n\ncompute_chiplet_area_mm2, MI210_core_breakdown_map, compute_total_die_map = (\n calc_compute_chiplet_area_mm2(arch_specs, verbose=True)\n)\nio_die_area_mm2, io_total_die_map = calc_io_die_area_mm2(arch_specs, verbose=True)\n\n\nvalues_mi210_sim = [\n compute_total_die_map[\"cores_area\"],\n compute_total_die_map[\"crossbar_area\"],\n io_total_die_map[\"global_buffer_area\"],\n io_total_die_map[\"mem_phy_area\"],\n io_total_die_map[\"mem_controller_area\"],\n io_total_die_map[\"device_phy_area\"],\n io_total_die_map[\"device_controller_area\"],\n 0,\n]\n\nplt.figure(figsize=(4, 2))\n\ncolors_matmul = sns.color_palette(\"flare_r\", 7)[5:6]\ncolors_normalization = sns.color_palette(\"summer\", 2)\ncolors_gelu = sns.color_palette(\"pink\", 5)[2:4]\ncolors_allreduce = sns.color_palette(\"Blues_r\", 2)\ncolors = (\n colors_matmul\n + colors_normalization\n + colors_gelu\n + colors_allreduce\n + sns.color_palette(\"Greys_r\", 1)\n)\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_a100)):\n plt.bar(1, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n bottom += value\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_a100_sim)):\n plt.bar(2, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_mi210)):\n plt.bar(3, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_mi210_sim)):\n plt.bar(4, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\nplt.ylabel(\"Area ($mm^2$)\")\nplt.xticks(\n [1, 2, 3, 4],\n [\"Real\\nGA100\", \"Simulated\\nGA100\", \"Real\\nAldebaran\", \"Simulated\\nAldebaran\"],\n)\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1, 1.1))\nplt.savefig(\"ae/figure6/figure6a.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)\n\n\nvalues_a100 = [3.75]\nvalues_mi210 = [4.02]\n\nvalues_a100_sim = [\n A100_core_breakdown_map[\"control_area\"],\n A100_core_breakdown_map[\"alu_area\"],\n A100_core_breakdown_map[\"sa_area\"],\n A100_core_breakdown_map[\"regfile_area\"],\n A100_core_breakdown_map[\"local_buffer_area\"],\n]\n\nvalues_mi210_sim = [\n MI210_core_breakdown_map[\"control_area\"],\n MI210_core_breakdown_map[\"alu_area\"],\n MI210_core_breakdown_map[\"sa_area\"],\n MI210_core_breakdown_map[\"regfile_area\"],\n MI210_core_breakdown_map[\"local_buffer_area\"],\n]\n\ncategories = [\n \"Control logic\",\n \"ALUs\",\n \"Systolic array\",\n \"Register file\",\n \"Local buffer\",\n]\n\ncolors = colors_matmul + colors_normalization + colors_allreduce\ncolor_gt = sns.color_palette(\"Greys_r\", 1)[0]\n\nplt.figure(figsize=(4, 1.5))\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_a100)):\n plt.bar(1, value, bottom=bottom, color=color_gt, width=0.5)\n bottom += value\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_a100_sim)):\n plt.bar(2, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n bottom += value\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_mi210)):\n plt.bar(3, value, bottom=bottom, color=color_gt, width=0.5)\n bottom += value\n\nbottom = 0\nfor i, (category, value) in enumerate(zip(categories, values_mi210_sim)):\n plt.bar(4, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\n\nplt.ylabel(\"Area ($mm^2$)\")\nplt.xticks(\n [1, 2, 3, 4],\n [\"Real\\nGA100\", \"Simulated\\nGA100\", \"Real\\nAldebaran\", \"Simulated\\nAldebaran\"],\n)\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1, 1))\nplt.savefig(\"ae/figure6/figure6b.pdf\", bbox_inches=\"tight\", pad_inches=0.01, dpi=300)\n" }, { "path": "ae/figure7/__init__.py", "content": "" }, { "path": "ae/figure7/change_core_size.py", "content": "from software_model.transformer import (\n TransformerBlockInitComputationTP,\n TransformerBlockAutoRegressionTP,\n)\nfrom software_model.utils import data_type_dict, Tensor\nfrom design_space_exploration.dse import template_to_system, read_architecture_template\nfrom multiprocessing import Process, Lock\nfrom cost_model.cost_model import calc_compute_chiplet_area_mm2, calc_io_die_area_mm2\nimport time\n\ninput_seq_length = 2048\nbatch_size = 8\noutput_seq_length = 1024\narch_specs = read_architecture_template(\"configs/GA100.json\")\ndevice_count = arch_specs[\"device_count\"]\n\nmodel_init = TransformerBlockInitComputationTP(\n d_model=12288,\n n_heads=96,\n device_count=device_count,\n data_type=data_type_dict[\"fp16\"],\n)\nmodel_auto_regression = TransformerBlockAutoRegressionTP(\n d_model=12288,\n n_heads=96,\n device_count=device_count,\n data_type=data_type_dict[\"fp16\"],\n)\n_ = model_init(\n Tensor([batch_size, input_seq_length, model_init.d_model], data_type_dict[\"fp16\"])\n)\n_ = model_auto_regression(\n Tensor([batch_size, 1, model_init.d_model], data_type_dict[\"fp16\"]),\n input_seq_length + output_seq_length,\n)\n\n\ndef test_core_size(core_configs, lock):\n name, core_count, sublane_count, array_width, vector_width, sram_KB = core_configs\n arch_specs[\"device\"][\"compute_chiplet\"][\"core_count\"] = core_count\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\"sublane_count\"] = sublane_count\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\"systolic_array\"][\n \"array_width\"\n ] = array_width\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\"systolic_array\"][\n \"array_height\"\n ] = array_width\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\"vector_unit\"][\n \"vector_width\"\n ] = vector_width\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\"SRAM_KB\"] = sram_KB\n # for area\n arch_specs[\"device\"][\"compute_chiplet\"][\"physical_core_count\"] = core_count\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\"vector_unit\"][\"int32_count\"] = (\n vector_width // 2\n )\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\"vector_unit\"][\"fp32_count\"] = (\n vector_width // 2\n )\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\"vector_unit\"][\"fp64_count\"] = (\n vector_width // 4\n )\n if vector_width <= 32:\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\"register_file\"][\n \"num_registers\"\n ] = (vector_width * 512)\n else:\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\"register_file\"][\n \"num_reg_files\"\n ] = (vector_width // 32)\n compute_area_mm2 = calc_compute_chiplet_area_mm2(arch_specs)\n io_area_mm2 = calc_io_die_area_mm2(arch_specs)\n print(f\"{name}, {compute_area_mm2}, {io_area_mm2}, {compute_area_mm2+io_area_mm2}\")\n # exit()\n system = template_to_system(arch_specs)\n auto_regression_latency_simulated = model_auto_regression.compile_and_simulate(\n system, \"heuristic-GPU\"\n )\n init_latency_simulated = model_init.compile_and_simulate(system, \"heuristic-GPU\")\n print(f\"{name}, {init_latency_simulated}, {auto_regression_latency_simulated}\")\n with lock:\n with open(f\"ae/figure7/core_size_results_init.csv\", \"a\") as f:\n f.write(\n f\"{name}, {compute_area_mm2+io_area_mm2}, {init_latency_simulated}, {model_init.simluate_log}\\n\"\n )\n with open(f\"ae/figure7/core_size_results_ar.csv\", \"a\") as f:\n f.write(\n f\"{name}, {compute_area_mm2+io_area_mm2}, {auto_regression_latency_simulated}, {model_auto_regression.simluate_log}\\n\"\n )\n\n\nlock = Lock()\nconfigs = [\n (\"A\", 128, 4, 8, 8, 192),\n (\"B\", 128, 4, 16, 32, 192),\n (\"C\", 128, 1, 32, 128, 192),\n (\"D\", 32, 1, 64, 512, 768),\n (\"E\", 8, 1, 128, 2048, 3072),\n]\n\nprocesses = [Process(target=test_core_size, args=(i, lock)) for i in configs]\n\ntry:\n for p in processes:\n p.start()\n\n while any(p.is_alive() for p in processes):\n time.sleep(1)\nexcept KeyboardInterrupt:\n print(\"Terminating processes...\")\n for p in processes:\n p.terminate()\n p.join()\n\n\nprint(\"All processes have finished.\")\n" }, { "path": "ae/figure7/plot_core_size.py", "content": "import matplotlib.pyplot as plt\nimport seaborn as sns\nimport csv\nimport pandas as pd\n\ncategories = [\n \"Q_K_V\",\n \"Q_mul_K\",\n \"A_mul_V\",\n \"Wo_proj\",\n \"W1_proj\",\n \"W2_proj\",\n \"Softmax\",\n \"LayerNorm_MHA\",\n \"LayerNorm_FFN\",\n \"GeLU\",\n \"AllReduce_MHA\",\n \"AllReduce_FFN\",\n]\ncol_names = [\"area\", \"latency\"] + categories\n\ncolors_matmul = sns.color_palette(\"flare_r\", 6)\ncolors_normalization = sns.color_palette(\"summer\", 3)\ncolors_gelu = sns.color_palette(\"pink\", 1)\ncolors_allreduce = sns.color_palette(\"Blues_r\", 2)\ncolors = colors_matmul + colors_normalization + colors_gelu + colors_allreduce\n\ncore_size_init = pd.read_csv(\n \"core_size_results_init.csv\", header=None, names=col_names, index_col=0\n)\ncore_size_init.index.astype(str)\ncore_size_ar = pd.read_csv(\n \"core_size_results_ar.csv\", header=None, names=col_names, index_col=0\n)\ncore_size_ar.index.astype(str)\n\n\ndf_sorted = core_size_init.sort_index()\nareas = df_sorted[\"area\"].tolist()\n# print(areas)\n# exit()\n# areas = [\n# 475.52039916931585,\n# 826.76355498007,\n# 826.76355498007,\n# 793.3380639020086,\n# 763.3465573533286,\n# ]\n\nplt.figure(figsize=(7, 3))\n\n# Create the stacked bar graph\nx = 0\nfor row_index in [\"A\", \"B\", \"C\", \"D\", \"E\"]:\n x = x + 1\n values = core_size_init.loc[row_index].tolist()\n bottom = 0\n for i, (category, value) in enumerate(zip(categories, values[2:])):\n if row_index == \"A\":\n plt.bar(x, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n else:\n plt.bar(x, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\n\nplt.ylabel(\"Latency (s)\")\nplt.xlabel(\"Configurations\")\nplt.xticks([1, 2, 3, 4, 5], [\"A\", \"B\", \"C\", \"D\", \"E\"])\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1.2, 1.05))\nplt.tight_layout()\nxticks = plt.gca().get_xticks()\nxticklabels = plt.gca().get_xticklabels()\nindex_to_color_red = list(xticks).index(2)\nxticklabels[index_to_color_red].set_color(\"#76B900\")\n\nax1 = plt.gca()\nax2 = ax1.twinx()\nax2.plot(\n [1, 2, 3, 4, 5],\n areas,\n color=\"dimgray\",\n linestyle=\"dashed\",\n marker=\"x\",\n label=\"Area\",\n)\nax2.set_ylabel(\"Area ($mm^2$)\")\nax2.set_ylim([0, 1000])\nplt.legend(loc=\"upper right\")\nplt.savefig(\n \"figure7a.pdf\",\n dpi=300,\n bbox_inches=\"tight\",\n pad_inches=0.01,\n)\nplt.show()\n\n\nplt.figure(figsize=(7, 3))\nx = 0\nfor row_index in [\"A\", \"B\", \"C\", \"D\", \"E\"]:\n x = x + 1\n values = core_size_ar.loc[row_index].tolist()\n bottom = 0\n for i, (category, value) in enumerate(zip(categories, values[2:])):\n value = value * 1e3\n if row_index == \"A\":\n plt.bar(x, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n else:\n plt.bar(x, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\n\n\n# Set the title, legend, and display the graph\n# plt.title(\n# \"Generation latency under different organization\"\n# )\nplt.ylabel(\"Latency (ms)\")\nplt.xlabel(\"Configurations\")\nplt.xticks([1, 2, 3, 4, 5], [\"A\", \"B\", \"C\", \"D\", \"E\"])\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1.2, 1.05))\nplt.tight_layout()\nxticks = plt.gca().get_xticks()\nxticklabels = plt.gca().get_xticklabels()\nindex_to_color_red = list(xticks).index(2)\nxticklabels[index_to_color_red].set_color(\"#76B900\")\nax1 = plt.gca()\nax2 = ax1.twinx()\nax2.plot(\n [1, 2, 3, 4, 5],\n areas,\n color=\"dimgrey\",\n linestyle=\"dashed\",\n marker=\"x\",\n label=\"Area\",\n)\nax2.set_ylabel(\"Area ($mm^2$)\")\nax2.set_ylim([0, 1000])\nplt.legend(loc=\"upper left\")\nplt.savefig(\n \"figure7b.pdf\",\n dpi=300,\n bbox_inches=\"tight\",\n pad_inches=0.01,\n)\n" }, { "path": "ae/figure7/run_figure7.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../..\n\npython -m ae.figure7.change_core_size\n\ncd ae/figure7\npython plot_core_size.py" }, { "path": "ae/figure8/__init__.py", "content": "" }, { "path": "ae/figure8/change_memory_bw.py", "content": "import json, re\nfrom hardware_model.compute_module import (\n VectorUnit,\n SystolicArray,\n Core,\n ComputeModule,\n overhead_dict,\n)\nfrom hardware_model.io_module import IOModule\nfrom hardware_model.memory_module import MemoryModule\nfrom hardware_model.device import Device\nfrom hardware_model.interconnect import LinkModule, InterConnectModule, TopologyType\nfrom hardware_model.system import System\nfrom software_model.transformer import (\n TransformerBlockInitComputationTP,\n TransformerBlockAutoRegressionTP,\n)\nfrom software_model.utils import data_type_dict, Tensor\nfrom cost_model.cost_model import calc_compute_chiplet_area_mm2, calc_io_die_area_mm2\nfrom math import ceil\n\nfrom design_space_exploration.dse import template_to_system, read_architecture_template\nfrom multiprocessing import Process, Lock\nimport time\nfrom cost_model.cost_model import calc_compute_chiplet_area_mm2, calc_io_die_area_mm2\n\n\ninput_seq_length = 2048\nbatch_size = 8\noutput_seq_length = 1024\narch_specs = read_architecture_template(\"configs/template.json\")\ndevice_count = arch_specs[\"device_count\"]\nmodel_init = TransformerBlockInitComputationTP(\n d_model=12288,\n n_heads=96,\n device_count=device_count,\n data_type=data_type_dict[\"fp16\"],\n)\nmodel_auto_regression = TransformerBlockAutoRegressionTP(\n d_model=12288,\n n_heads=96,\n device_count=device_count,\n data_type=data_type_dict[\"fp16\"],\n)\n_ = model_init(\n Tensor([batch_size, input_seq_length, model_init.d_model], data_type_dict[\"fp16\"])\n)\n_ = model_auto_regression(\n Tensor([batch_size, 1, model_init.d_model], data_type_dict[\"fp16\"]),\n input_seq_length + output_seq_length,\n)\n\n\ndef test_memory_bandwidth(memory_bandwidth, lock):\n arch_specs[\"device\"][\"io\"][\"memory_channel_physical_count\"] = memory_bandwidth\n arch_specs[\"device\"][\"io\"][\"memory_channel_active_count\"] = memory_bandwidth\n compute_area_mm2 = calc_compute_chiplet_area_mm2(arch_specs)\n io_area_mm2 = calc_io_die_area_mm2(arch_specs)\n print(\n f\"{memory_bandwidth}, {compute_area_mm2}, {io_area_mm2}, {compute_area_mm2+io_area_mm2}\"\n )\n system = template_to_system(arch_specs)\n auto_regression_latency_simulated = model_auto_regression.compile_and_simulate(\n system, \"heuristic-GPU\"\n )\n init_latency_simulated = model_init.compile_and_simulate(system, \"heuristic-GPU\")\n print(\n f\"{memory_bandwidth}, {init_latency_simulated}, {auto_regression_latency_simulated}\"\n )\n with lock:\n with open(f\"ae/figure8/memory_bw_results_bs{batch_size}_init.csv\", \"a\") as f:\n f.write(\n f\"{memory_bandwidth*400}, {compute_area_mm2+io_area_mm2}, {init_latency_simulated}, {model_init.simluate_log}\\n\"\n )\n with open(f\"ae/figure8/memory_bw_results_bs{batch_size}_ar.csv\", \"a\") as f:\n f.write(\n f\"{memory_bandwidth*400}, {compute_area_mm2+io_area_mm2}, {auto_regression_latency_simulated}, {model_auto_regression.simluate_log}\\n\"\n )\n\n\nlock = Lock()\nprocesses = [\n Process(target=test_memory_bandwidth, args=(i, lock))\n for i in [1, 2, 3, 4, 5, 6, 7, 8]\n]\n\ntry:\n for p in processes:\n p.start()\n\n while any(p.is_alive() for p in processes):\n time.sleep(1)\nexcept KeyboardInterrupt:\n print(\"Terminating processes...\")\n for p in processes:\n p.terminate()\n p.join()\n\n\nprint(\"All processes have finished.\")\n" }, { "path": "ae/figure8/plot_memory_bw.py", "content": "import matplotlib.pyplot as plt\nimport seaborn as sns\nimport csv\nimport pandas as pd\n\ncategories = [\n \"Q_K_V\",\n \"Q_mul_K\",\n \"A_mul_V\",\n \"Wo_proj\",\n \"W1_proj\",\n \"W2_proj\",\n \"Softmax\",\n \"LayerNorm_MHA\",\n \"LayerNorm_FFN\",\n \"GeLU\",\n \"AllReduce_MHA\",\n \"AllReduce_FFN\",\n]\ncol_names = [\"area\", \"latency\"] + categories\n\ncolors_matmul = sns.color_palette(\"flare_r\", 6)\ncolors_normalization = sns.color_palette(\"summer\", 3)\ncolors_gelu = sns.color_palette(\"pink\", 1)\ncolors_allreduce = sns.color_palette(\"Blues_r\", 2)\ncolors = colors_matmul + colors_normalization + colors_gelu + colors_allreduce\n\nbatch_size = 8\n\nresults_init = pd.read_csv(\n f\"memory_bw_results_bs{batch_size}_init.csv\",\n header=None,\n names=col_names,\n index_col=0,\n)\nresults_init.index.astype(int)\nresults_ar = pd.read_csv(\n f\"memory_bw_results_bs{batch_size}_ar.csv\",\n header=None,\n names=col_names,\n index_col=0,\n)\nresults_ar.index.astype(int)\n\n\nplt.figure(figsize=(7, 3))\n\n# Create the stacked bar graph\nx = 0\nx_labels = [i * 400 for i in [1, 2, 3, 4, 5, 6, 7, 8]]\nfor row_index in x_labels:\n x = x + 1\n values = results_init.loc[row_index].tolist()\n bottom = 0\n for i, (category, value) in enumerate(zip(categories, values[2:])):\n if row_index == x_labels[0]:\n plt.bar(x, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n else:\n plt.bar(x, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\n\n\n# Set the title, legend, and display the graph\n# plt.title(\n# \"Prefilling Latency per Layer\"\n# )\nplt.ylabel(\"Latency (s)\")\nplt.xlabel(\"Memory bandwidth (GB/s)\")\nplt.xticks([1, 2, 3, 4, 5, 6, 7, 8], x_labels)\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1, 1.05))\nplt.tight_layout()\nxticks = plt.gca().get_xticks()\nxticklabels = plt.gca().get_xticklabels()\nindex_to_color_red = list(xticks).index(5)\nxticklabels[index_to_color_red].set_color(\"red\")\nplt.savefig(f\"figure8a.pdf\", dpi=300, bbox_inches=\"tight\", pad_inches=0.01)\nplt.show()\n\n\nplt.figure(figsize=(7, 3))\nx = 0\nfor row_index in x_labels:\n x = x + 1\n values = results_ar.loc[row_index].tolist()\n bottom = 0\n for i, (category, value) in enumerate(zip(categories, values[2:])):\n value = value * 1e3\n if row_index == x_labels[0]:\n plt.bar(x, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n else:\n plt.bar(x, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\n\n\n# Set the title, legend, and display the graph\n# plt.title(\n# \"Generation Latency per Layer per Token\"\n# )\nplt.ylabel(\"Latency (ms)\")\nplt.xlabel(\"Memory bandwidth (GB/s)\")\nplt.xticks([1, 2, 3, 4, 5, 6, 7, 8], x_labels)\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1, 1.05))\nplt.tight_layout()\nxticks = plt.gca().get_xticks()\nxticklabels = plt.gca().get_xticklabels()\nindex_to_color_red = list(xticks).index(5)\nxticklabels[index_to_color_red].set_color(\"red\")\nplt.savefig(f\"figure8b.pdf\", dpi=300, bbox_inches=\"tight\", pad_inches=0.01)\n" }, { "path": "ae/figure8/run_figure8.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../..\n\npython -m ae.figure8.change_memory_bw\n\ncd ae/figure8\npython plot_memory_bw.py" }, { "path": "ae/figure9/__init__.py", "content": "" }, { "path": "ae/figure9/change_l1_cache.py", "content": "from software_model.transformer import (\n TransformerBlockInitComputationTP,\n TransformerBlockAutoRegressionTP,\n)\nfrom software_model.utils import data_type_dict, Tensor\nfrom design_space_exploration.dse import template_to_system, read_architecture_template\nfrom multiprocessing import Process, Lock\nfrom cost_model.cost_model import calc_compute_chiplet_area_mm2, calc_io_die_area_mm2\nimport time\n\ninput_seq_length = 2048\nbatch_size = 8\noutput_seq_length = 1024\narch_specs = read_architecture_template(\"configs/template.json\")\ndevice_count = arch_specs[\"device_count\"]\n\nmodel_init = TransformerBlockInitComputationTP(\n d_model=12288,\n n_heads=96,\n device_count=device_count,\n data_type=data_type_dict[\"fp16\"],\n)\nmodel_auto_regression = TransformerBlockAutoRegressionTP(\n d_model=12288,\n n_heads=96,\n device_count=device_count,\n data_type=data_type_dict[\"fp16\"],\n)\n_ = model_init(\n Tensor([batch_size, input_seq_length, model_init.d_model], data_type_dict[\"fp16\"])\n)\n_ = model_auto_regression(\n Tensor([batch_size, 1, model_init.d_model], data_type_dict[\"fp16\"]),\n input_seq_length + output_seq_length,\n)\n\n\ndef test_SRAM_KB(SRAM_KB, lock):\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\"SRAM_KB\"] = SRAM_KB\n compute_area_mm2 = calc_compute_chiplet_area_mm2(arch_specs)\n io_area_mm2 = calc_io_die_area_mm2(arch_specs)\n print(\n f\"{SRAM_KB}, {compute_area_mm2}, {io_area_mm2}, {compute_area_mm2+io_area_mm2}\"\n )\n system = template_to_system(arch_specs)\n auto_regression_latency_simulated = model_auto_regression.compile_and_simulate(\n system, \"heuristic-GPU\"\n )\n init_latency_simulated = model_init.compile_and_simulate(system, \"heuristic-GPU\")\n print(f\"{SRAM_KB}, {init_latency_simulated}, {auto_regression_latency_simulated}\")\n with lock:\n with open(f\"ae/figure9/l1_cache_results_init.csv\", \"a\") as f:\n f.write(\n f\"{SRAM_KB}, {compute_area_mm2+io_area_mm2}, {init_latency_simulated}, {model_init.simluate_log}\\n\"\n )\n with open(f\"ae/figure9/l1_cache_results_ar.csv\", \"a\") as f:\n f.write(\n f\"{SRAM_KB}, {compute_area_mm2+io_area_mm2}, {auto_regression_latency_simulated}, {model_auto_regression.simluate_log}\\n\"\n )\n\n\n# for SRAM_KB in [64, 128, 192, 256, 512, 1024]:\n# test_SRAM_KB(SRAM_KB, None)\n\nlock = Lock()\nprocesses = [\n Process(target=test_SRAM_KB, args=(i, lock)) for i in [64, 128, 192, 256, 512, 1024]\n]\n\ntry:\n for p in processes:\n p.start()\n\n while any(p.is_alive() for p in processes):\n time.sleep(1)\nexcept KeyboardInterrupt:\n print(\"Terminating processes...\")\n for p in processes:\n p.terminate()\n p.join()\n\n\nprint(\"All processes have finished.\")\n\n# for SRAM_KB in [64, 128, 192, 256, 512, 1024]:\n# arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\n# \"SRAM_KB\"\n# ] = SRAM_KB\n# system=template_to_system(arch_specs)\n# auto_regression_latency_simulated = model_auto_regression.compile_and_simulate(system, 'heuristic-GPU')\n# init_latency_simulated = model_init.compile_and_simulate(system, 'heuristic-GPU')\n# print(f'{SRAM_KB}, {init_latency_simulated}, {auto_regression_latency_simulated}')\n# with open(f'test/case_study/l1_cache/l1_cache_results.csv', 'a') as f:\n# f.write(f'{SRAM_KB}, {init_latency_simulated}, {auto_regression_latency_simulated}\\n')\n" }, { "path": "ae/figure9/plot_l1_cache.py", "content": "import matplotlib.pyplot as plt\nimport seaborn as sns\nimport csv\nimport pandas as pd\n\ncategories = [\n \"Q_K_V\",\n \"Q_mul_K\",\n \"A_mul_V\",\n \"Wo_proj\",\n \"W1_proj\",\n \"W2_proj\",\n \"Softmax\",\n \"LayerNorm_MHA\",\n \"LayerNorm_FFN\",\n \"GeLU\",\n \"AllReduce_MHA\",\n \"AllReduce_FFN\",\n]\ncol_names = [\"area\", \"latency\"] + categories\n\ncolors_matmul = sns.color_palette(\"flare_r\", 6)\ncolors_normalization = sns.color_palette(\"summer\", 3)\ncolors_gelu = sns.color_palette(\"pink\", 1)\ncolors_allreduce = sns.color_palette(\"Blues_r\", 2)\ncolors = colors_matmul + colors_normalization + colors_gelu + colors_allreduce\n\nresults_init = pd.read_csv(\n \"l1_cache_results_init.csv\", header=None, names=col_names, index_col=0\n)\nresults_init.index.astype(int)\nresults_ar = pd.read_csv(\n \"l1_cache_results_ar.csv\", header=None, names=col_names, index_col=0\n)\nresults_ar.index.astype(int)\n\nareas = [\n 782.1048032068737,\n 794.1065561553206,\n 826.76355498007,\n 848.4527315580167,\n 913.304090096728,\n 1064.9121549472263,\n]\n\nplt.figure(figsize=(7, 3))\n\n# Create the stacked bar graph\nx = 0\nx_labels = [64, 128, 192, 256, 512, 1024]\nfor row_index in x_labels:\n x = x + 1\n values = results_init.loc[row_index].tolist()\n bottom = 0\n for i, (category, value) in enumerate(zip(categories, values[2:])):\n if row_index == x_labels[0]:\n plt.bar(x, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n else:\n plt.bar(x, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\n\n\n# Set the title, legend, and display the graph\n# plt.title(\n# \"Prefilling Latency per Layer\"\n# )\nplt.ylabel(\"Latency (s)\")\nplt.xlabel(\"Local buffer size (KB)\")\nplt.xticks([1, 2, 3, 4, 5, 6], x_labels)\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1.2, 1.03))\nplt.tight_layout()\nxticks = plt.gca().get_xticks()\nxticklabels = plt.gca().get_xticklabels()\nindex_to_color_red = list(xticks).index(3)\nxticklabels[index_to_color_red].set_color(\"#76B900\")\n\nax1 = plt.gca()\nax2 = ax1.twinx()\nax2.plot(\n [1, 2, 3, 4, 5, 6],\n areas,\n color=\"dimgray\",\n linestyle=\"dashed\",\n marker=\"x\",\n label=\"Area\",\n)\nax2.set_ylabel(\"Area ($mm^2$)\")\nax2.set_ylim([0, 1200])\nplt.legend(loc=\"upper center\")\n\nplt.savefig(\"figure9a.pdf\", dpi=300, bbox_inches=\"tight\", pad_inches=0.01)\nplt.show()\n\n\nplt.figure(figsize=(7, 3))\nx = 0\nfor row_index in x_labels:\n x = x + 1\n values = results_ar.loc[row_index].tolist()\n bottom = 0\n for i, (category, value) in enumerate(zip(categories, values[2:])):\n value = value * 1e3\n if row_index == x_labels[0]:\n plt.bar(x, value, bottom=bottom, color=colors[i], label=category, width=0.5)\n else:\n plt.bar(x, value, bottom=bottom, color=colors[i], width=0.5)\n bottom += value\n\n\n# Set the title, legend, and display the graph\n# plt.title(\n# \"Generation Latency per Layer per Token\"\n# )\nplt.ylabel(\"Latency (ms)\")\nplt.xlabel(\"Local buffer size (KB)\")\nplt.xticks([1, 2, 3, 4, 5, 6], x_labels)\nhandles, labels = plt.gca().get_legend_handles_labels()\nplt.legend(handles[::-1], labels[::-1], loc=\"upper left\", bbox_to_anchor=(1.2, 1.03))\nplt.tight_layout()\nxticks = plt.gca().get_xticks()\nxticklabels = plt.gca().get_xticklabels()\nindex_to_color_red = list(xticks).index(3)\nxticklabels[index_to_color_red].set_color(\"#76B900\")\nax1 = plt.gca()\nax1.set_ylim([0, 1.2])\nax2 = ax1.twinx()\nax2.plot(\n [1, 2, 3, 4, 5, 6],\n areas,\n color=\"dimgray\",\n linestyle=\"dashed\",\n marker=\"x\",\n label=\"Area\",\n)\nax2.set_ylabel(\"Area ($mm^2$)\")\nax2.set_ylim([0, 1200])\nplt.legend(loc=\"upper center\")\nplt.savefig(\"figure9b.pdf\", dpi=300, bbox_inches=\"tight\", pad_inches=0.01)\n" }, { "path": "ae/figure9/run_figure9.sh", "content": "rm *.csv\nrm *.pdf\n\ncd ../..\n\npython -m ae.figure9.change_l1_cache\n\ncd ae/figure9\npython plot_l1_cache.py" }, { "path": "configs/GA100.json", "content": "{\n \"name\": \"NVIDIA A100(80GB)x4\",\n \"device_count\": 4,\n \"interconnect\": {\n \"link\": {\n \"name\": \"NVLink3\",\n \"bandwidth_per_direction_byte\": 25e9,\n \"bandwidth_both_directions_byte\": 50e9,\n \"latency_second\": 8.92e-6,\n \"flit_size_byte\": 16,\n \"header_size_byte\": 16,\n \"max_payload_size_byte\": 256\n },\n \"link_count_per_device\": 12,\n \"topology\": \"FC\"\n },\n \"device\": {\n \"frequency_Hz\": 1410e6,\n \"compute_chiplet_count\": 1,\n \"compute_chiplet\": {\n \"physical_core_count\": 128,\n \"core_count\": 128,\n \"process_node\": \"7nm\",\n \"core\": {\n \"sublane_count\": 4,\n \"systolic_array\": {\n \"array_width\": 16,\n \"array_height\": 16,\n \"data_type\": \"fp16\",\n \"mac_per_cycle\": 1\n },\n \"vector_unit\": {\n \"vector_width\": 32,\n \"flop_per_cycle\": 4,\n \"data_type\": \"fp16\",\n \"int32_count\": 16,\n \"fp16_count\": 0,\n \"fp32_count\": 16,\n \"fp64_count\": 8\n },\n \"register_file\": {\n \"num_reg_files\": 1,\n \"num_registers\": 16384,\n \"register_bitwidth\":32,\n \"num_rdwr_ports\":4\n },\n \"SRAM_KB\": 192\n }\n },\n \"memory_protocol\": \"HBM2e\",\n \"_memory_protocol_list\": [\n \"HBM2e\",\n \"DDR4\",\n \"DDR5\",\n \"PCIe4\",\n \"PCIe5\"\n ],\n \"io\": {\n \"process_node\": \"7nm\",\n \"global_buffer_MB\": 48,\n \"physical_global_buffer_MB\": 48,\n \"global_buffer_bandwidth_per_cycle_byte\": 5120,\n \"memory_channel_physical_count\": 6,\n \"memory_channel_active_count\": 5,\n \"pin_count_per_channel\": 1024,\n \"bandwidth_per_pin_bit\": 3.2e9\n },\n \"memory\": {\n \"total_capacity_GB\": 80\n }\n }\n}" }, { "path": "configs/ga102_template.json", "content": "{\n \"name\": \"NVIDIA A100(80GB)x4\",\n \"device_count\": 4,\n \"interconnect\": {\n \"link\": {\n \"name\": \"NVLink3\",\n \"bandwidth_per_direction_byte\": 25e9,\n \"bandwidth_both_directions_byte\": 50e9,\n \"latency_second\": 8.92e-6,\n \"flit_size_byte\": 16,\n \"header_size_byte\": 16,\n \"max_payload_size_byte\": 256\n },\n \"link_count_per_device\": 12,\n \"topology\": \"FC\"\n },\n \"device\": {\n \"frequency_Hz\": 1410e6,\n \"compute_chiplet_count\": 1,\n \"compute_chiplet\": {\n \"physical_core_count\": 128,\n \"core_count\": 108,\n \"process_node\": \"7nm\",\n \"core\": {\n \"sublane_count\": 4,\n \"systolic_array\": {\n \"array_width\": 16,\n \"array_height\": 16,\n \"data_type\": \"fp16\",\n \"mac_per_cycle\": 1\n },\n \"vector_unit\": {\n \"vector_width\": 32,\n \"flop_per_cycle\": 4,\n \"data_type\": \"fp16\",\n \"int32_count\": 16,\n \"fp16_count\": 0,\n \"fp32_count\": 32,\n \"fp64_count\": 0.5\n },\n \"SRAM_KB\": 128\n }\n },\n \"memory_protocol\": \"HBM2e\",\n \"_memory_protocol_list\": [\n \"HBM2e\",\n \"DDR4\",\n \"DDR5\",\n \"PCIe4\",\n \"PCIe5\"\n ],\n \"io\": {\n \"process_node\": \"7nm\",\n \"global_buffer_MB\": 40,\n \"physical_global_buffer_MB\": 48,\n \"global_buffer_bandwidth_per_cycle_byte\": 5120,\n \"memory_channel_physical_count\": 6,\n \"memory_channel_active_count\": 5,\n \"pin_count_per_channel\": 1024,\n \"bandwidth_per_pin_bit\": 3.2e9\n },\n \"memory\": {\n \"total_capacity_GB\": 80\n }\n }\n}" }, { "path": "configs/generation_system.json", "content": "{\n \"name\": \"NVIDIA A100(80GB)x4\",\n \"device_count\": 4,\n \"interconnect\": {\n \"link\": {\n \"name\": \"NVLink3\",\n \"bandwidth_per_direction_byte\": 25e9,\n \"bandwidth_both_directions_byte\": 50e9,\n \"latency_second\": 8.92e-6,\n \"flit_size_byte\": 16,\n \"header_size_byte\": 16,\n \"max_payload_size_byte\": 256\n },\n \"link_count_per_device\": 12,\n \"topology\": \"FC\"\n },\n \"device\": {\n \"frequency_Hz\": 1410e6,\n \"compute_chiplet_count\": 1,\n \"compute_chiplet\": {\n \"physical_core_count\": 128,\n \"core_count\": 128,\n \"process_node\": \"7nm\",\n \"core\": {\n \"sublane_count\": 4,\n \"systolic_array\": {\n \"array_width\": 8,\n \"array_height\": 8,\n \"data_type\": \"fp16\",\n \"mac_per_cycle\": 1\n },\n \"vector_unit\": {\n \"vector_width\": 8,\n \"flop_per_cycle\": 4,\n \"data_type\": \"fp16\",\n \"int32_count\": 4,\n \"fp16_count\": 32,\n \"fp32_count\": 0,\n \"fp64_count\": 0\n },\n \"register_file\": {\n \"num_reg_files\": 1,\n \"num_registers\": 4096,\n \"register_bitwidth\": 32,\n \"num_rdwr_ports\": 4\n },\n \"SRAM_KB\": 48\n }\n },\n \"memory_protocol\": \"HBM2e\",\n \"_memory_protocol_list\": [\n \"HBM2e\",\n \"DDR4\",\n \"DDR5\",\n \"PCIe4\",\n \"PCIe5\"\n ],\n \"io\": {\n \"process_node\": \"7nm\",\n \"global_buffer_MB\": 24,\n \"physical_global_buffer_MB\": 24,\n \"global_buffer_bandwidth_per_cycle_byte\": 2560,\n \"memory_channel_physical_count\": 6,\n \"memory_channel_active_count\": 5,\n \"pin_count_per_channel\": 1024,\n \"bandwidth_per_pin_bit\": 3.2e9\n },\n \"memory\": {\n \"total_capacity_GB\": 80\n }\n }\n}" }, { "path": "configs/latency_design.json", "content": "{\n \"name\": \"NVIDIA A100(80GB)x4\",\n \"device_count\": 4,\n \"interconnect\": {\n \"link\": {\n \"name\": \"NVLink3\",\n \"bandwidth_per_direction_byte\": 25e9,\n \"bandwidth_both_directions_byte\": 50e9,\n \"latency_second\": 8.92e-6,\n \"flit_size_byte\": 16,\n \"header_size_byte\": 16,\n \"max_payload_size_byte\": 256\n },\n \"link_count_per_device\": 12,\n \"topology\": \"FC\"\n },\n \"device\": {\n \"frequency_Hz\": 1410e6,\n \"compute_chiplet_count\": 1,\n \"compute_chiplet\": {\n \"physical_core_count\": 64,\n \"core_count\": 64,\n \"process_node\": \"7nm\",\n \"core\": {\n \"sublane_count\": 4,\n \"systolic_array\": {\n \"array_width\": 16,\n \"array_height\": 16,\n \"data_type\": \"fp16\",\n \"mac_per_cycle\": 1\n },\n \"vector_unit\": {\n \"vector_width\": 32,\n \"flop_per_cycle\": 4,\n \"data_type\": \"fp16\",\n \"int32_count\": 16,\n \"fp16_count\": 0,\n \"fp32_count\": 16,\n \"fp64_count\": 8\n },\n \"register_file\": {\n \"num_reg_files\": 1,\n \"num_registers\": 16384,\n \"register_bitwidth\": 32,\n \"num_rdwr_ports\": 4\n },\n \"SRAM_KB\": 192\n }\n },\n \"memory_protocol\": \"HBM2e\",\n \"_memory_protocol_list\": [\n \"HBM2e\",\n \"DDR4\",\n \"DDR5\",\n \"PCIe4\",\n \"PCIe5\"\n ],\n \"io\": {\n \"process_node\": \"7nm\",\n \"global_buffer_MB\": 24,\n \"physical_global_buffer_MB\": 24,\n \"global_buffer_bandwidth_per_cycle_byte\": 2560,\n \"memory_channel_physical_count\": 6,\n \"memory_channel_active_count\": 5,\n \"pin_count_per_channel\": 1024,\n \"bandwidth_per_pin_bit\": 3.2e9\n },\n \"memory\": {\n \"total_capacity_GB\": 80\n }\n }\n}" }, { "path": "configs/mi210.json", "content": "{\n \"name\": \"AMD MI210\",\n \"device_count\": 4,\n \"interconnect\": {\n \"link\": {\n \"name\": \"NVLink3\",\n \"bandwidth_per_direction_byte\": 25e9,\n \"bandwidth_both_directions_byte\": 50e9,\n \"latency_second\": 8.92e-6,\n \"flit_size_byte\": 16,\n \"header_size_byte\": 16,\n \"max_payload_size_byte\": 256\n },\n \"link_count_per_device\": 12,\n \"topology\": \"FC\"\n },\n \"device\": {\n \"frequency_Hz\": 1400e6,\n \"compute_chiplet_count\": 1,\n \"compute_chiplet\": {\n \"core_count\": 104,\n \"process_node\": \"7nm\",\n \"core\": {\n \"sublane_count\": 4,\n \"systolic_array\": {\n \"array_width\": 16,\n \"array_height\": 16,\n \"data_type\": \"fp16\",\n \"mac_per_cycle\": 1\n },\n \"vector_unit\": {\n \"vector_width\": 16,\n \"flop_per_cycle\": 2,\n \"data_type\": \"fp32\",\n \"int32_count\": 16,\n \"fp32_count\": 16,\n \"fp64_count\": 8\n },\n \"SRAM_KB\": 128\n }\n },\n \"memory_protocol\": \"HBM2e\",\n \"_memory_protocol_list\": [\n \"HBM2e\",\n \"DDR4\",\n \"DDR5\",\n \"PCIe4\",\n \"PCIe5\"\n ],\n \"io\": {\n \"process_node\": \"7nm\",\n \"global_buffer_MB\": 8,\n \"global_buffer_bandwidth_per_cycle_byte\": 4096,\n \"memory_channel_physical_count\": 6,\n \"memory_channel_active_count\": 4,\n \"pin_count_per_channel\": 1024,\n \"bandwidth_per_pin_bit\": 3.2e9\n },\n \"memory\": {\n \"total_capacity_GB\": 64\n }\n }\n}" }, { "path": "configs/mi210_template.json", "content": "{\n \"name\": \"AMD MI210\",\n \"device_count\": 4,\n \"interconnect\": {\n \"link\": {\n \"name\": \"InfinityFabric\",\n \"bandwidth_per_direction_byte\": 25e9,\n \"bandwidth_both_directions_byte\": 50e9,\n \"latency_second\": 8.92e-6,\n \"flit_size_byte\": 16,\n \"header_size_byte\": 16,\n \"max_payload_size_byte\": 256\n },\n \"link_count_per_device\": 8,\n \"topology\": \"FC\"\n },\n \"device\": {\n \"frequency_Hz\": 1410e6,\n \"compute_chiplet_count\": 1,\n \"compute_chiplet\": {\n \"physical_core_count\": 112,\n \"core_count\": 108,\n \"process_node\": \"6nm\",\n \"core\": {\n \"sublane_count\": 4,\n \"systolic_array\": {\n \"array_width\": 16,\n \"array_height\": 16,\n \"data_type\": \"fp16\",\n \"mac_per_cycle\": 1\n },\n \"vector_unit\": {\n \"vector_width\": 16,\n \"flop_per_cycle\": 4,\n \"data_type\": \"fp16\",\n \"int32_count\": 16,\n \"fp16_count\": 0,\n \"fp32_count\": 0,\n \"fp64_count\": 16\n },\n \"register_file\": {\n \"num_reg_files\": 64,\n \"num_registers\": 512,\n \"register_bitwidth\":32,\n \"num_rdwr_ports\":4\n },\n \"SRAM_KB\": 80\n }\n },\n \"memory_protocol\": \"HBM2e\",\n \"_memory_protocol_list\": [\n \"HBM2e\",\n \"DDR4\",\n \"DDR5\",\n \"PCIe4\",\n \"PCIe5\"\n ],\n \"io\": {\n \"process_node\": \"6nm\",\n \"global_buffer_MB\": 8,\n \"physical_global_buffer_MB\": 8,\n \"global_buffer_bandwidth_per_cycle_byte\": 5120,\n \"memory_channel_physical_count\": 8,\n \"memory_channel_active_count\": 8,\n \"pin_count_per_channel\": 1024,\n \"bandwidth_per_pin_bit\": 3.2e9\n },\n \"memory\": {\n \"total_capacity_GB\": 80\n }\n }\n}" }, { "path": "configs/prefilling_system.json", "content": "{\n \"name\": \"NVIDIA A100(80GB)x4\",\n \"device_count\": 4,\n \"interconnect\": {\n \"link\": {\n \"name\": \"NVLink3\",\n \"bandwidth_per_direction_byte\": 25e9,\n \"bandwidth_both_directions_byte\": 50e9,\n \"latency_second\": 8.92e-6,\n \"flit_size_byte\": 16,\n \"header_size_byte\": 16,\n \"max_payload_size_byte\": 256\n },\n \"link_count_per_device\": 12,\n \"topology\": \"FC\"\n },\n \"device\": {\n \"frequency_Hz\": 1410e6,\n \"compute_chiplet_count\": 1,\n \"compute_chiplet\": {\n \"physical_core_count\": 64,\n \"core_count\": 64,\n \"process_node\": \"7nm\",\n \"core\": {\n \"sublane_count\": 4,\n \"systolic_array\": {\n \"array_width\": 32,\n \"array_height\": 32,\n \"data_type\": \"fp16\",\n \"mac_per_cycle\": 1\n },\n \"vector_unit\": {\n \"vector_width\": 32,\n \"flop_per_cycle\": 4,\n \"data_type\": \"fp16\",\n \"int32_count\": 16,\n \"fp16_count\": 0,\n \"fp32_count\": 16,\n \"fp64_count\": 8\n },\n \"register_file\": {\n \"num_reg_files\": 1,\n \"num_registers\": 16384,\n \"register_bitwidth\": 32,\n \"num_rdwr_ports\": 4\n },\n \"SRAM_KB\": 768\n }\n },\n \"memory_protocol\": \"PCIe5\",\n \"_memory_protocol_list\": [\n \"HBM2e\",\n \"DDR4\",\n \"DDR5\",\n \"PCIe4\",\n \"PCIe5\"\n ],\n \"io\": {\n \"process_node\": \"7nm\",\n \"global_buffer_MB\": 48,\n \"physical_global_buffer_MB\": 48,\n \"global_buffer_bandwidth_per_cycle_byte\": 5120,\n \"memory_channel_physical_count\": 256,\n \"memory_channel_active_count\": 256,\n \"pin_count_per_channel\": 1,\n \"bandwidth_per_pin_bit\": 32e9\n },\n \"memory\": {\n \"total_capacity_GB\": 160\n }\n }\n}" }, { "path": "configs/template.json", "content": "{\n \"name\": \"NVIDIA A100(80GB)x4\",\n \"device_count\": 4,\n \"interconnect\": {\n \"link\": {\n \"name\": \"NVLink3\",\n \"bandwidth_per_direction_byte\": 25e9,\n \"bandwidth_both_directions_byte\": 50e9,\n \"latency_second\": 8.92e-6,\n \"flit_size_byte\": 16,\n \"header_size_byte\": 16,\n \"max_payload_size_byte\": 256\n },\n \"link_count_per_device\": 12,\n \"topology\": \"FC\"\n },\n \"device\": {\n \"frequency_Hz\": 1410e6,\n \"compute_chiplet_count\": 1,\n \"compute_chiplet\": {\n \"physical_core_count\": 128,\n \"core_count\": 108,\n \"process_node\": \"7nm\",\n \"core\": {\n \"sublane_count\": 4,\n \"systolic_array\": {\n \"array_width\": 16,\n \"array_height\": 16,\n \"data_type\": \"fp16\",\n \"mac_per_cycle\": 1\n },\n \"vector_unit\": {\n \"vector_width\": 32,\n \"flop_per_cycle\": 4,\n \"data_type\": \"fp16\",\n \"int32_count\": 16,\n \"fp16_count\": 0,\n \"fp32_count\": 16,\n \"fp64_count\": 8\n },\n \"register_file\": {\n \"num_reg_files\": 1,\n \"num_registers\": 16384,\n \"register_bitwidth\":32,\n \"num_rdwr_ports\":4\n },\n \"SRAM_KB\": 192\n }\n },\n \"memory_protocol\": \"HBM2e\",\n \"_memory_protocol_list\": [\n \"HBM2e\",\n \"DDR4\",\n \"DDR5\",\n \"PCIe4\",\n \"PCIe5\"\n ],\n \"io\": {\n \"process_node\": \"7nm\",\n \"global_buffer_MB\": 40,\n \"physical_global_buffer_MB\": 48,\n \"global_buffer_bandwidth_per_cycle_byte\": 5120,\n \"memory_channel_physical_count\": 6,\n \"memory_channel_active_count\": 5,\n \"pin_count_per_channel\": 1024,\n \"bandwidth_per_pin_bit\": 3.2e9\n },\n \"memory\": {\n \"total_capacity_GB\": 80\n }\n }\n}" }, { "path": "cost_model/__init__.py", "content": "## made this for first commit\n" }, { "path": "cost_model/cost_examples.py", "content": "import cost_model.cost_model as cost_model\nimport json\n\n# example chip with a 32 wide vector, 16x16 SA, 256kb cache core, 8 cores per die\n# io die with 64 mb cache, 8 nvlinks, 32 pcie phys\n# all at 5nm and 7nm\n\nwith open(\"./configs/prefilling_system.json\", \"r\") as f:\n # with open('../configs/mi210_template.json', 'r') as f:\n configs_dict = json.load(f)\n\n# print(configs_dict['device'])\n# print(data['device']['compute_chiplet_count'])\n\ncompute_area = cost_model.calc_compute_chiplet_area_mm2(configs_dict)\nio_area = cost_model.calc_io_die_area_mm2(configs_dict)\nprint(\n f\"compute area: {compute_area}, io area: {io_area}, total area: {compute_area+io_area}\"\n)\n\nexit(0)\ncore_compute_area_mm2 = cost_model.calc_compute_core_area_mm2(\n 32,\n 16,\n 2**18,\n cost_model.transistor_density_7nm,\n cost_model.sram_bit_cell_density_7nm,\n)\nio_die_area_mm2 = cost_model.calc_io_die_area_mm2(\n 2**25,\n cost_model.PCIE5,\n 32,\n 8,\n cost_model.transistor_density_7nm,\n cost_model.sram_bit_cell_density_7nm,\n)\nprint(core_compute_area_mm2)\nprint(io_die_area_mm2)\n" }, { "path": "cost_model/cost_model.py", "content": "# Author: August Ning aning@princeton.edu\n# Date started: 12 October 2023\n# This file is the cost model for Naivesim\n\nimport numpy as np\nimport math\n# import supply_chain.supply_chain_model as scm\nimport cost_model.supply_chain.supply_chain_model as scm\n\n# lots of parameters required for calculating silicon die area cost\n\n# these are in terms of million transistors per mm2\ntransistor_density_7nm = scm.transistor_density_arr[scm.PN_7_INDEX]\ntransistor_density_6nm = 114.2\ntransistor_density_5nm = scm.transistor_density_arr[scm.PN_5_INDEX]\n\nsram_bit_cell_density_7nm = 1.70e-07\nsram_bit_cell_density_6nm = 1.40e-07\nsram_bit_cell_density_5nm = 1.25e-07\n\n# cache size overheads derived from cacti for cache sizes\n# 4096, 8192, 16384, ..., 1 MB\ncache_area_efficiency_arr = [0.076, 0.142, 0.247, 0.393, 0.559, \\\n 0.704, 0.526, 0.602, 0.561]\n\n# fpu transistor counts are for 64 bit FPU, based off Ariane and OpenPiton's SPARC T1\n# assume that fp32 are half the transistors\n# int32 transistor count is based off of Ariane's Mult and OpenPiton's SPARC T1\n# systolic array is for 1x1 area\n\n# scale FPU area by mantissa bits quadratically\nfpu64_transistor_count = 685300\nfpu32_transistor_count = fpu64_transistor_count * ((23 / 52) ** 2)\nfpu16_transistor_count = fpu64_transistor_count * ((10 / 52) ** 2)\n\nint32_transistor_count = 177690\n\n# based off of A100 SM and MI 210 CU\n# these overheads are per sublane, per vector width\n# (ex 32 for A100, 16 for MI 210)\nper_sublane_control_transistor_count = 996200\nnvidia_per_sublane_control_transistor_count = 725650\namd_per_sublane_control_transistor_count = 1534500\n\nper_sublane_control_dict = {'nvidia':per_sublane_control_transistor_count, \\\n 'amd':per_sublane_control_transistor_count}\n\nper_core_comm_transistor_count = 44300000\nnvidia_per_core_comm_transistor_count = 55000000\namd_per_core_comm_transistor_count = 33600000\n\nper_core_comm_dict = {'nvidia':per_core_comm_transistor_count, \\\n 'amd':per_core_comm_transistor_count}\n\n\n# memory controllers scale with process node, but PHYs do not\n# pcie, ddr, hbm\n# note: DDR link unit is 32 bits\npcie5_phy_mm2_per_lane = 0.64\npcie4_phy_mm2_per_lane = 0.48\nddr5_phy_mm2_per_link_unit = 1.45\nhbm2e_phy_mm2_per_link_unit = 10.45\nnvlink3_phy_mm2_per_link_unit = 1.888\nnvlink4_phy_mm2_per_link_unit = 0.965\ninfinity_fabric_phy_mm2_per_link_unit = 5.69\n\npcie5_ctrl_transistors_per_lane = 5372100\npcie4_ctrl_transistors_per_lane = 3962500\nddr5_ctrl_transistors_per_link_unit = 90446400\nhbm2e_ctrl_transistors_per_link_unit = 552743000\nnvlink3_ctrl_transistors_per_link_unit = 74632000\nnvlink4_ctrl_transistors_per_link_unit = 86628000\ninfinity_fabric_ctrl_transistors_per_link_unit = 348148000\n\n# mem tech keywords\nPCIE5 = 'PCIe5'\nPCIE4 = 'PCIe4'\nDDR5 = 'DDR5'\nHBM = 'HBM2e'\nNVLINK3 = 'NVLink3'\nNVLINK4 = 'NVLink4'\nINFINITYFABRIC = 'InfinityFabric'\n\n# average via dramexchange spot price, Oct 2023\nddr5_cost_per_gb = 2.4\nhbm_cost_per_gb = 7\n\n# return die area for a dimension x dimension SA with a\n# give bitwidth FPU at a given process node\n# right now, we model each PE's MAC as a FPU\ndef calc_systolic_array_area_mm2(dimension_x, dimension_y, bitwidth, transistor_density_mil_mm2):\n if bitwidth == 'fp64':\n total_transistor_count = fpu64_transistor_count * dimension_x * dimension_y\n elif bitwidth == 'fp32':\n total_transistor_count = fpu32_transistor_count * dimension_x * dimension_y\n elif bitwidth == 'fp16':\n total_transistor_count = fpu16_transistor_count * dimension_x * dimension_y\n\n return total_transistor_count / 1e6 / transistor_density_mil_mm2\n\n# vector width corresponds to number of FPUs you have\ndef calc_vector_area_mm2(int32_count, fp16_count, fp32_count, fp64_count, transistor_density_mil_mm2):\n total_transistor_count = 0\n total_transistor_count += int32_count * int32_transistor_count\n total_transistor_count += fp16_count * fpu16_transistor_count\n total_transistor_count += fp32_count * fpu32_transistor_count\n total_transistor_count += fp64_count * fpu64_transistor_count\n \n return total_transistor_count / 1e6 / transistor_density_mil_mm2\n\n# for cache designs, if the desired capacity is larger than the max cache unit\n# split them up into multiple units of the max capacity\n# min cache size is 4096 bytes\ndef calc_cache_sram_area_mm2(capacity_bytes, sram_bitcell_area_mm2, max_cache_unit_bytes=(2**19)):\n if capacity_bytes > max_cache_unit_bytes:\n num_cache_units = math.ceil(capacity_bytes / max_cache_unit_bytes)\n unit_size_bytes = max_cache_unit_bytes\n else:\n num_cache_units = 1\n unit_size_bytes = capacity_bytes\n\n # cache size model is for capacity of 4096 bytes to 1 MB\n if unit_size_bytes < 2 ** 12:\n unit_size_bytes = 2 ** 12\n\n area_efficiency_index = math.ceil(math.log(unit_size_bytes, 2)) - 12\n area_efficiency_factor = cache_area_efficiency_arr[area_efficiency_index]\n unit_cache_area = unit_size_bytes * 8 * sram_bitcell_area_mm2 / area_efficiency_factor\n cache_area = num_cache_units * unit_cache_area\n return cache_area\n\n# area model comes from EMPIRE\n# num_reg_files: how many distinct register files each sublanes has\n# D: how many registers there are in each RF\n# W: bits per register\n# P: number of read/write ports\ndef calc_reg_file_area(num_reg_files, D, W, P, transistor_density_mil_mm2):\n area_90nm_um2 = (3.29 * 10**4) - (1.09 * 10**3 * D) - (8.83 * 10**2 * W) - (5.55 * 10**3 * P) \\\n + (5.35 * 10**1 * D * W) + (1.50 * 10**-2 * D**2) + (1.08 * 10**-2 * W**2) \\\n + (5.86 * 10**-1 * P**2) + (1.42 * 10**2 * D * P) + (3.68 * 10**2 * W * P)\n \n # need to convert um2 to mm2, convert to 7nm\n area_90nm_mm2 = area_90nm_um2 / 1e6\n area_mm2 = area_90nm_mm2 * (scm.transistor_density_arr[scm.PN_90_INDEX] / transistor_density_mil_mm2)\n total_reg_file_area = num_reg_files * area_mm2\n return total_reg_file_area\n\n# for width, for PCIe and NVLink, it is the whole lane\n# for DDR and HBM, it's 128 bits and 1024 bits respectively\ndef calc_mem_controller_area_mm2(mem_tech, width, transistor_density_mil_mm2):\n controller_transistor_count = -1\n\n if mem_tech == PCIE5:\n controller_transistor_count = pcie5_ctrl_transistors_per_lane * width\n elif mem_tech == PCIE4:\n controller_transistor_count = pcie4_ctrl_transistors_per_lane * width\n elif mem_tech == DDR5:\n controller_transistor_count = ddr5_ctrl_transistors_per_link_unit * width\n elif mem_tech == HBM:\n controller_transistor_count = hbm2e_ctrl_transistors_per_link_unit * width\n elif mem_tech == NVLINK3:\n controller_transistor_count = nvlink3_ctrl_transistors_per_link_unit * width\n elif mem_tech == NVLINK4:\n controller_transistor_count = nvlink4_ctrl_transistors_per_link_unit * width\n elif mem_tech == INFINITYFABRIC:\n controller_transistor_count = infinity_fabric_ctrl_transistors_per_link_unit * width\n\n return (controller_transistor_count / 1e6) / transistor_density_mil_mm2\n\ndef calc_mem_phy_area_mm2(mem_tech, width):\n if mem_tech == PCIE5:\n return pcie5_phy_mm2_per_lane * width\n elif mem_tech == PCIE4:\n return pcie4_phy_mm2_per_lane * width\n elif mem_tech == DDR5:\n return ddr5_phy_mm2_per_link_unit * width\n elif mem_tech == HBM:\n return hbm2e_phy_mm2_per_link_unit * width\n elif mem_tech == NVLINK3:\n return nvlink3_phy_mm2_per_link_unit * width\n elif mem_tech == NVLINK4:\n return nvlink4_phy_mm2_per_link_unit * width\n elif mem_tech == INFINITYFABRIC:\n return infinity_fabric_phy_mm2_per_link_unit * width\n else:\n return -1\n\ndef find_logic_sram_transistor_density(process_node):\n if '7' in process_node:\n return transistor_density_7nm, sram_bit_cell_density_7nm\n elif '6' in process_node:\n return transistor_density_6nm, sram_bit_cell_density_6nm\n elif '5' in process_node:\n return transistor_density_5nm, sram_bit_cell_density_5nm\n\n raise Exception(\"Invalid Process Node\")\n\n\n# a compute core consists of a fixed control overhead\n# a specified width fp32 vector engine\n# a specified dimmension fp16 systolic array\n# a specified L1 cache\n# at a specified process node\n# NB: you can fit multiple cores onto a single die for chiplet systems\ndef calc_compute_chiplet_area_mm2(configs_dict, verbose=False):\n total_die_map = {}\n core_breakdown_map = {}\n device_name = configs_dict['name']\n device_brand = 'nvidia' if 'nvidia' in device_name.lower() else 'amd'\n vector_width = configs_dict['device']['compute_chiplet']['core']['vector_unit']['vector_width']\n vector_int32_count = configs_dict['device']['compute_chiplet']['core']['vector_unit']['int32_count']\n vector_fp16_count = configs_dict['device']['compute_chiplet']['core']['vector_unit']['fp16_count']\n vector_fp32_count = configs_dict['device']['compute_chiplet']['core']['vector_unit']['fp32_count']\n vector_fp64_count = configs_dict['device']['compute_chiplet']['core']['vector_unit']['fp64_count']\n sa_dim_x = configs_dict['device']['compute_chiplet']['core']['systolic_array']['array_width']\n sa_dim_y = configs_dict['device']['compute_chiplet']['core']['systolic_array']['array_height']\n sa_bitwidth = configs_dict['device']['compute_chiplet']['core']['systolic_array']['data_type']\n num_reg_files = configs_dict['device']['compute_chiplet']['core']['register_file']['num_reg_files']\n num_registers = configs_dict['device']['compute_chiplet']['core']['register_file']['num_registers']\n register_bitwidth = configs_dict['device']['compute_chiplet']['core']['register_file']['register_bitwidth']\n num_rdwr_ports = configs_dict['device']['compute_chiplet']['core']['register_file']['num_rdwr_ports']\n sublane_count = configs_dict['device']['compute_chiplet']['core']['sublane_count']\n cache_size_bytes = configs_dict['device']['compute_chiplet']['core']['SRAM_KB'] * (2 ** 10)\n process_node = configs_dict['device']['compute_chiplet']['process_node']\n cores_per_chiplet = configs_dict['device']['compute_chiplet']['physical_core_count']\n\n # each sublane has a SA and vector unit. a core is made up of sublanes. a chiplet has multiple cores\n transistor_density_mil_mm2, sram_density_bitcell_mm2 = find_logic_sram_transistor_density(process_node)\n per_sublane_area_mm2 = 0\n per_sublane_control_area_mm2 = per_sublane_control_dict[device_brand] / 1e6 / transistor_density_mil_mm2\n per_sublane_area_mm2 += (vector_width * per_sublane_control_area_mm2)\n control_logic_area = per_sublane_area_mm2 * sublane_count\n\n per_lane_vector_area = calc_vector_area_mm2(vector_int32_count, vector_fp16_count, vector_fp32_count, vector_fp64_count, transistor_density_mil_mm2)\n per_sublane_area_mm2 += per_lane_vector_area\n\n per_lane_sa_area = calc_systolic_array_area_mm2(sa_dim_x, sa_dim_y, sa_bitwidth, transistor_density_mil_mm2)\n per_sublane_area_mm2 += per_lane_sa_area\n\n per_lane_regfile_area = calc_reg_file_area(num_reg_files, num_registers, register_bitwidth, num_rdwr_ports, transistor_density_mil_mm2)\n per_sublane_area_mm2 += per_lane_regfile_area\n\n per_core_compute_area_mm2 = per_sublane_area_mm2 * sublane_count\n cache_area_mm2 = calc_cache_sram_area_mm2(cache_size_bytes, sram_density_bitcell_mm2)\n per_core_area_mm2 = per_core_compute_area_mm2 + cache_area_mm2\n\n core_breakdown_map['total_core_area'] = per_core_area_mm2\n core_breakdown_map['control_area'] = control_logic_area\n core_breakdown_map['alu_area'] = per_lane_vector_area * sublane_count\n core_breakdown_map['sa_area'] = per_lane_sa_area * sublane_count\n core_breakdown_map['regfile_area'] = per_lane_regfile_area * sublane_count\n core_breakdown_map['local_buffer_area'] = cache_area_mm2\n\n total_cores_area = per_core_area_mm2 * cores_per_chiplet\n total_crossbar_area = (per_core_comm_dict[device_brand] / 1e6 / transistor_density_mil_mm2) * cores_per_chiplet\n # each core has an area overhead to connect to the xbar\n compute_chiplet_area_mm2 = total_cores_area + total_crossbar_area\n\n total_die_map['total_area'] = compute_chiplet_area_mm2\n total_die_map['cores_area'] = total_cores_area\n total_die_map['crossbar_area'] = total_crossbar_area\n if verbose:\n return compute_chiplet_area_mm2, core_breakdown_map, total_die_map\n else:\n return compute_chiplet_area_mm2\n\n\n# NB: for mem_tech, if you are using DDR or HBM, it will be 128 bits and 1024 bits respectively per lane\n# for PCIe and NVLink, specify the number of lanes (128bits per lane)\n# def calc_io_die_area_mm2(cache_size_bytes, mem_tech, mem_tech_width, num_nvlink_phys, \\\n# transistor_density_mil_mm2, sram_density_bitcell_mm2):\ndef calc_io_die_area_mm2(config_dict, verbose=False):\n total_die_map = {}\n\n cache_size_bytes = config_dict['device']['io']['physical_global_buffer_MB'] * (2 ** 20)\n mem_tech = config_dict['device']['memory_protocol']\n num_mem_tech_units = config_dict['device']['io']['memory_channel_physical_count']\n gpu_gpu_comm_tech = config_dict['interconnect']['link']['name']\n num_gpu_gpu_comm_phy = config_dict['interconnect']['link_count_per_device']\n process_node = config_dict['device']['io']['process_node']\n\n transistor_density_mil_mm2, sram_density_bitcell_mm2 = find_logic_sram_transistor_density(process_node)\n io_die_area_mm2 = 0\n io_die_area_mm2 += calc_cache_sram_area_mm2(cache_size_bytes, sram_density_bitcell_mm2)\n global_buffer_area = io_die_area_mm2\n\n # mem tech for communicating to off chip memory\n mem_phy_area = calc_mem_phy_area_mm2(mem_tech, num_mem_tech_units)\n mem_controller_area = calc_mem_controller_area_mm2(mem_tech, num_mem_tech_units, transistor_density_mil_mm2)\n io_die_area_mm2 += mem_phy_area\n io_die_area_mm2 += mem_controller_area\n\n # every IO die has a few NV links for chip to chip communication\n device_phy_area = calc_mem_phy_area_mm2(gpu_gpu_comm_tech, num_gpu_gpu_comm_phy)\n device_controller_area = calc_mem_controller_area_mm2(gpu_gpu_comm_tech, num_gpu_gpu_comm_phy, transistor_density_mil_mm2)\n io_die_area_mm2 += device_phy_area\n io_die_area_mm2 += device_controller_area\n\n total_die_map['total_die_area'] = io_die_area_mm2\n total_die_map['global_buffer_area'] = global_buffer_area\n total_die_map['mem_phy_area'] = mem_phy_area\n total_die_map['mem_controller_area'] = mem_controller_area\n total_die_map['device_phy_area'] = device_phy_area\n total_die_map['device_controller_area'] = device_controller_area\n\n if verbose:\n return io_die_area_mm2, total_die_map\n else:\n return io_die_area_mm2\n" }, { "path": "cost_model/regfile_area.py", "content": "def calculate_regfile_area(D, W, P):\n area_90nm_um2 = (3.29 * 10**4) - (1.09 * 10**3 * D) - (8.83 * 10**2 * W) - (5.55 * 10**3 * P) \\\n + (5.35 * 10**1 * D * W) + (1.50 * 10**-2 * D**2) + (1.08 * 10**-2 * W**2) \\\n + (5.86 * 10**-1 * P**2) + (1.42 * 10**2 * D * P) + (3.68 * 10**2 * W * P)\n \n # need to convert um2 to mm2, convert to 7nm\n area_90nm_mm2 = area_90nm_um2 / 1e6\n area_7nm_mm2 = area_90nm_mm2 * (1.6 / 96.3)\n return area_7nm_mm2\n\nreg_area = calculate_regfile_area(16384, 32, 4)\nprint(reg_area)\n\nreg_area = 64 * calculate_regfile_area(512, 32, 4)\nprint(reg_area)\n\nreg_area = calculate_regfile_area(800, 32, 4)\nprint(reg_area)" }, { "path": "design_space_exploration/__init__.py", "content": "" }, { "path": "design_space_exploration/dse.py", "content": "import json, re\nfrom hardware_model.compute_module import (\n VectorUnit,\n SystolicArray,\n Core,\n ComputeModule,\n overhead_dict,\n)\nfrom hardware_model.io_module import IOModule\nfrom hardware_model.memory_module import MemoryModule\nfrom hardware_model.device import Device\nfrom hardware_model.interconnect import LinkModule, InterConnectModule, TopologyType\nfrom hardware_model.system import System\nfrom software_model.transformer import (\n TransformerBlockInitComputationTP,\n TransformerBlockAutoRegressionTP,\n)\nfrom software_model.utils import data_type_dict, Tensor\n# from cost_model.cost_model import calc_compute_chiplet_area_mm2, calc_io_die_area_mm2\nfrom math import ceil\n\ndef read_architecture_template(file_path):\n with open(file_path, \"r\") as f:\n arch_specs = json.load(f)\n return arch_specs\n\n\ndef template_to_system(arch_specs):\n device_specs = arch_specs[\"device\"]\n compute_chiplet_specs = device_specs[\"compute_chiplet\"]\n io_specs = device_specs[\"io\"]\n core_specs = compute_chiplet_specs[\"core\"]\n sublane_count = core_specs[\"sublane_count\"]\n # vector unit\n vector_unit_specs = core_specs[\"vector_unit\"]\n vector_unit = VectorUnit(\n sublane_count\n * vector_unit_specs[\"vector_width\"]\n * vector_unit_specs[\"flop_per_cycle\"],\n int(re.search(r\"(\\d+)\", vector_unit_specs[\"data_type\"]).group(1)) // 8,\n 35,\n vector_unit_specs[\"vector_width\"],\n sublane_count,\n )\n # systolic array\n systolic_array_specs = core_specs[\"systolic_array\"]\n systolic_array = SystolicArray(\n systolic_array_specs[\"array_height\"],\n systolic_array_specs[\"array_width\"],\n systolic_array_specs[\"mac_per_cycle\"],\n int(re.search(r\"(\\d+)\", systolic_array_specs[\"data_type\"]).group(1)) // 8,\n int(re.search(r\"(\\d+)\", systolic_array_specs[\"data_type\"]).group(1)) // 8,\n )\n # core\n core = Core(\n vector_unit,\n systolic_array,\n sublane_count,\n core_specs[\"SRAM_KB\"] * 1024,\n )\n # compute module\n compute_module = ComputeModule(\n core,\n compute_chiplet_specs[\"core_count\"] * device_specs[\"compute_chiplet_count\"],\n device_specs[\"frequency_Hz\"],\n io_specs[\"global_buffer_MB\"] * 1024 * 1024,\n io_specs[\"global_buffer_bandwidth_per_cycle_byte\"],\n overhead_dict[\"A100\"],\n )\n # io module\n io_module = IOModule(\n io_specs[\"memory_channel_active_count\"]\n * io_specs[\"pin_count_per_channel\"]\n * io_specs[\"bandwidth_per_pin_bit\"]\n // 8,\n 1e-6,\n )\n # memory module\n memory_module = MemoryModule(\n device_specs[\"memory\"][\"total_capacity_GB\"] * 1024 * 1024 * 1024\n )\n # device\n device = Device(compute_module, io_module, memory_module)\n # interconnect\n interconnect_specs = arch_specs[\"interconnect\"]\n link_specs = interconnect_specs[\"link\"]\n link_module = LinkModule(\n link_specs[\"bandwidth_per_direction_byte\"],\n link_specs[\"bandwidth_both_directions_byte\"],\n link_specs[\"latency_second\"],\n link_specs[\"flit_size_byte\"],\n link_specs[\"max_payload_size_byte\"],\n link_specs[\"header_size_byte\"],\n )\n interconnect_module = InterConnectModule(\n arch_specs[\"device_count\"],\n TopologyType.FC\n if interconnect_specs[\"topology\"] == \"FC\"\n else TopologyType.RING,\n link_module,\n interconnect_specs[\"link_count_per_device\"],\n )\n\n # system\n system = System(device, interconnect_module)\n\n return system\n\n\ndef test_template_to_system():\n arch_specs = read_architecture_template(\"configs/template.json\")\n A100_system = template_to_system(arch_specs)\n bs = 8\n s = 2048\n model = TransformerBlockInitComputationTP(\n d_model=12288,\n n_heads=96,\n device_count=4,\n data_type=data_type_dict[\"fp16\"],\n )\n _ = model(Tensor([bs, s, 12288], data_type_dict[\"fp16\"]))\n model.roofline_model(A100_system)\n\n\ndef find_cheapest_design(\n d_model,\n n_heads,\n n_layers,\n batch_size,\n input_seq_length,\n init_latency,\n output_seq_length,\n auto_regression_latency,\n \n):\n i=0\n smallest_total_area_mm2=float('inf')\n best_arch_specs=None\n arch_specs = read_architecture_template(\"configs/template.json\")\n for device_count in [4, 8, 12, 16]:\n model_init = TransformerBlockInitComputationTP(\n d_model=12288,\n n_heads=96,\n device_count=device_count,\n data_type=data_type_dict[\"fp16\"],\n )\n model_auto_regression = TransformerBlockAutoRegressionTP(\n d_model=12288,\n n_heads=96,\n device_count=device_count,\n data_type=data_type_dict[\"fp16\"],)\n _ = model_init(Tensor([batch_size, input_seq_length, model_init.d_model], data_type_dict[\"fp16\"]))\n _ = model_auto_regression(Tensor([batch_size, 1, model_init.d_model],data_type_dict[\"fp16\"]), input_seq_length+output_seq_length)\n arch_specs[\"device_count\"] = device_count\n if device_count <= 4:\n topology = \"FC\"\n else:\n topology = \"RING\"\n arch_specs[\"interconnect\"][\"topology\"] = topology\n for link_count_per_device in [6, 12, 18, 24]:\n arch_specs[\"interconnect\"][\"link_count_per_device\"] = link_count_per_device\n # device\n for core_count in [32, 64, 128, 256]:\n arch_specs[\"device\"][\"compute_chiplet\"][\"core_count\"] = core_count\n # core\n for sublane_count in [1, 2, 4, 8]:\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\n \"sublane_count\"\n ] = sublane_count\n # systolic array\n for array_height in [16, 32, 64, 128]:\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\n \"systolic_array\"\n ][\"array_height\"] = array_height\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\n \"systolic_array\"\n ][\"array_width\"] = array_height\n # vector unit\n for vector_width in [16, 32, 64, 128]:\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\n \"vector_unit\"\n ][\"vector_width\"] = vector_width\n for SRAM_KB in [64, 128, 256, 512, 1024]:\n arch_specs[\"device\"][\"compute_chiplet\"][\"core\"][\n \"SRAM_KB\"\n ] = SRAM_KB\n # global buffer\n for total_global_buffer_MB in [\n 80,\n 160,\n 240,\n 320,\n 400,\n 480,\n 640,\n 800,\n 960,\n ]:\n global_buffer_MB = (\n total_global_buffer_MB // device_count\n )\n global_buffer_bandwidth_per_cycle_byte = (\n 5120 * global_buffer_MB // 40\n )\n arch_specs[\"device\"][\"io\"][\n \"global_buffer_MB\"\n ] = global_buffer_MB\n arch_specs[\"device\"][\"io\"][\n \"global_buffer_bandwidth_per_cycle_byte\"\n ] = global_buffer_bandwidth_per_cycle_byte\n # memory\n memory_capacity_requirement_GB = ceil(model_auto_regression.memory_requirement*n_layers/1e9/16)*16\n # print(f\"memory_capacity_requirement_GB={model_auto_regression.memory_requirement*n_layers/1e9}\")\n # exit()\n for memory_protocol in [\n \"HBM2e\",\n \"DDR5\",\n \"PCIe5\",\n # \"GDDR6X\"\n ]:\n arch_specs['device']['memory_protocol']=memory_protocol\n if memory_protocol == \"HBM2e\":\n # 400 GB/s per channel, 16 GB\n channel_count=memory_capacity_requirement_GB // 16\n if channel_count>8:\n continue\n channel_count_list = [channel_count]\n pin_count_per_channel=1024\n bandwidth_per_pin_bit=3.2e9\n elif memory_protocol == \"DDR5\":\n # 19.2 GB/s per channel, 2 channel per dimm\n channel_count_list = [16, 24, 32]\n pin_count_per_channel=32\n bandwidth_per_pin_bit=4.8e9\n elif memory_protocol == \"PCIe5\":\n # 4 GB/s per channel\n channel_count_list = [64, 96, 128]\n pin_count_per_channel=1\n bandwidth_per_pin_bit=32e9\n # elif memory_protocol == \"GDDR6X\":\n # # 84 GB/s per channel, 2 GB\n # channel_count_list= memo\n for channel_count in channel_count_list:\n arch_specs['device']['memory']['total_capacity_GB'] = memory_capacity_requirement_GB\n arch_specs['device']['io']['memory_channel_active_count'] = channel_count\n arch_specs['device']['io']['memory_channel_physical_count'] = channel_count\n arch_specs['device']['io']['pin_count_per_channel'] = pin_count_per_channel\n arch_specs['device']['io']['bandwidth_per_pin_bit'] = bandwidth_per_pin_bit\n \n total_area_mm2=calc_compute_chiplet_area_mm2(arch_specs)+calc_io_die_area_mm2(arch_specs)\n # print(f\"channel_count={arch_specs['device']['io']['memory_channel_active_count']},total area={total_area_mm2}\")\n if total_area_mm2>900:\n continue\n system=template_to_system(arch_specs)\n init_roofline_latency=model_init.roofline_model(system)*n_layers\n if init_roofline_latency>init_latency:\n continue\n \n auto_regression_roofline_latency=model_auto_regression.roofline_model(system)*n_layers\n if auto_regression_roofline_latency>auto_regression_latency:\n continue\n auto_regression_latency_simulated = model_auto_regression.compile_and_simulate(system, 'heuristic-GPU')\n if auto_regression_latency_simulated>auto_regression_latency:\n continue\n init_latency_simulated = model_init.compile_and_simulate(system, 'heuristic-GPU')\n if init_latency_simulated>init_latency:\n continue\n if total_area_mm2*device_count