{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Google - Fast or Slow? - Detailed EDA\n\n## About this Competition\nMachine learning compilers play an important role in translating ML programs (high-level) into the optimized code (low-level) which can be executed on the specific hardware. In this competition, tensor computational graphs (represented as **H**igh **L**evel **O**prerations (**HLO**)), their configurations and the corresponding execution times are given. The competitors are challenged to predict the **runtime of each graph-configuration pair** (*i.e.,* the execution time of the graph when compiled with the specific configuration). Actually, it can be formulated as a **ranking problem**, and our goal is to **rank configurations of a given graph**.\n\n## About this Notebook\nIn this kernel, I will go through the provided raw data, trying to clarify confusion in my mind when I first step into this competition. Also, a naive baseline model is implemented to illustrate how the evaluation metric works.\n\n<a id=\"toc\"></a>\n## Table of Contents\n* [1. Data Appearance](#data_appearance)\n* [2. ML Compiler Optimizations](#ml_optim)\n    * *[Graph-Level Optimizations](#graph_lv)*\n    * *[Kernel-Level Optimizations](#kernel_lv)*\n* [3. Basic Graph Statistics](#graph_stats)\n    * *[layout-nlp-default](#lnd)*\n    * *[layout-nlp-random](#lnr)*\n    * *[layout-xla-default](#lnd)*\n    * *[layout-xla-random](#lnr)*\n    * *[tile-xla](#tx)*\n* [4. Our Target - Runtime](#runtime)\n* More analysis is coming soon...\n* [Reference](#ref)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"markdown","source":"## Import Packages","metadata":{}},{"cell_type":"code","source":"import os\nimport random\nimport warnings\nfrom collections import defaultdict\nfrom pathlib import Path\nfrom tqdm.notebook import tqdm\nfrom typing import Any, Dict, List, Optional, Union, Tuple\n\nimport networkx as nx\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport matplotlib.gridspec as gridspec\nimport matplotlib.image as img\nimport seaborn as sns\nimport plotly.graph_objects as go\nfrom matplotlib.axes import Axes\nfrom scipy.stats import mode, skew, kurtosis\nfrom seaborn import JointGrid\n\n# Configuration\nwarnings.simplefilter('ignore')\nsns.set_style(\"darkgrid\")\ncolors = sns.color_palette('Set2')\nrandom.seed(42)","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-09-09T05:01:49.737712Z","iopub.execute_input":"2023-09-09T05:01:49.738199Z","iopub.status.idle":"2023-09-09T05:01:51.315604Z","shell.execute_reply.started":"2023-09-09T05:01:49.738150Z","shell.execute_reply":"2023-09-09T05:01:51.314495Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define Paths, Metadata and Utilities","metadata":{}},{"cell_type":"code","source":"# Paths\nDATA_ROOT = Path(\"/kaggle/input/predict-ai-model-runtime/npz_all/npz\")\n\n# Metadata\n# ==Data==\n# Compiler optimization\nOPTIM = [\"layout\", \"tile\"]\n# Source\nSRC = [\"xla\", \"nlp\"]\n# Search strategy\nSEARCH = [\"default\", \"random\"]\n# Dataset split\nSPLIT = [\"train\", \"valid\", \"test\"]\n# Collection\nCOLL = [\n    \"layout-nlp-default\",\n    \"layout-nlp-random\",\n    \"layout-xla-default\",\n    \"layout-xla-random\",\n    \"tile-xla\"\n]","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:10:40.574577Z","iopub.execute_input":"2023-09-04T09:10:40.575225Z","iopub.status.idle":"2023-09-04T09:10:40.583373Z","shell.execute_reply.started":"2023-09-04T09:10:40.575192Z","shell.execute_reply":"2023-09-04T09:10:40.581341Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DataChecker(object):\n    \"\"\"Basic data checker.\n    \n    Parameters:\n        optim: compiler optimization, the choices are as follows:\n            {\"layout\", \"tile\"}\n    \"\"\"\n\n    CHECK_ITEMS: List[str] = [\"col1_deprecated\", \"op_align\", \"n_cfg_nodes\"]\n    \n    def __init__(self, optim: str) -> None:\n        self.optim = optim\n\n        # Optimization-specific setup\n        self._optim_root = DATA_ROOT/optim\n        self._src = SRC if optim == \"layout\" else [\"xla\"]\n        self._search = SEARCH if optim == \"layout\" else [\"\"]\n        \n    def check(self) -> Dict[str, List[str]]:\n        \"\"\"Run basic data checking.\n        \n        Return:\n            npz_illeg: .npz files failing specific check items\n        \"\"\"\n        npz_illeg = defaultdict(list)\n        for src in self._src:\n            for search in self._search:\n                for split in SPLIT:\n                    data_dir = self._optim_root/src/search/split\n        \n                    print(f\"Check {data_dir}...\")\n                    for file in os.listdir(data_dir):\n                        if not file.endswith(\".npz\"): continue\n                        \n                        data_file = str(data_dir/file)\n                        npz_tmp = dict(np.load(data_file))\n                        if not self._check_col1_deprecated(npz_tmp[\"node_feat\"]):\n                            npz_illeg[\"col1_deprecated\"].append(data_file)\n                        if not self._check_op_align(npz_tmp[\"node_feat\"], npz_tmp[\"node_opcode\"]):\n                            npz_illeg[\"op_align\"].append(data_file)\n                        if (self.optim == \"layout\"\n                            and not self._check_n_cfg_nodes(npz_tmp[\"node_feat\"], npz_tmp[\"node_config_ids\"])):\n                            npz_illeg[\"n_cfg_nodes\"].append(data_file)\n\n        return npz_illeg\n\n    def _check_col1_deprecated(self, node_feat: np.ndarray) -> bool:\n        \"\"\"Check whether column 1 of node_feat is deprecated for all files.\n        \n        Parameters:\n            node_feat: node feature array\n\n        Return:\n            pass_check: if True, the checking is passed\n        \"\"\"\n        col1_deprecated = np.all(node_feat[:, 1] == 0)\n        pass_check = True if col1_deprecated else False\n\n        return pass_check\n\n    def _check_op_align(self, node_feat: np.ndarray, node_opcode: np.ndarray) -> bool:\n        \"\"\"Check whether each node has its own op-code.\n        \n        Parameters:\n            node_feat: node feature array\n            node_opcode: node op-codes\n    \n        Return:\n            pass_check: if True, the checking is passed\n        \"\"\"\n        op_align = len(node_opcode) == len(node_feat)\n        pass_check = True if op_align else False\n    \n        return pass_check\n\n    def _check_n_cfg_nodes(self, node_feat: np.ndarray, node_config_id: np.ndarray) -> bool:\n        \"\"\"Check whether the number of configurable nodes is always less\n        than the total number of nodes.\n    \n        Check if `nc` <= `n`.\n    \n        Parameters:\n            node_feat: node feature array\n            node_config_id: configurable node identifier\n    \n        Return:\n            pass_check: if True, the checking is passed\n        \"\"\"\n        n_cfg_nodes = len(node_config_id) <= len(node_feat)\n        pass_check = True if n_cfg_nodes else False\n    \n        return pass_check\n    \n    \nclass SeabornFig2Grid():\n    \"\"\"See https://stackoverflow.com/questions/35042255/.\"\"\"\n\n    def __init__(self, seaborngrid, fig, subplot_spec):\n        self.fig = fig\n        self.sg = seaborngrid\n        self.subplot = subplot_spec\n        if isinstance(self.sg, sns.axisgrid.FacetGrid) or \\\n            isinstance(self.sg, sns.axisgrid.PairGrid):\n            self._movegrid()\n        elif isinstance(self.sg, sns.axisgrid.JointGrid):\n            self._movejointgrid()\n        self._finalize()\n\n    def _movegrid(self):\n        \"\"\" Move PairGrid or Facetgrid \"\"\"\n        self._resize()\n        n = self.sg.axes.shape[0]\n        m = self.sg.axes.shape[1]\n        self.subgrid = gridspec.GridSpecFromSubplotSpec(n,m, subplot_spec=self.subplot)\n        for i in range(n):\n            for j in range(m):\n                self._moveaxes(self.sg.axes[i,j], self.subgrid[i,j])\n\n    def _movejointgrid(self):\n        \"\"\" Move Jointgrid \"\"\"\n        h= self.sg.ax_joint.get_position().height\n        h2= self.sg.ax_marg_x.get_position().height\n        r = int(np.round(h/h2))\n        self._resize()\n        self.subgrid = gridspec.GridSpecFromSubplotSpec(r+1,r+1, subplot_spec=self.subplot)\n\n        self._moveaxes(self.sg.ax_joint, self.subgrid[1:, :-1])\n        self._moveaxes(self.sg.ax_marg_x, self.subgrid[0, :-1])\n        self._moveaxes(self.sg.ax_marg_y, self.subgrid[1:, -1])\n\n    def _moveaxes(self, ax, gs):\n        #https://stackoverflow.com/a/46906599/4124317\n        ax.remove()\n        ax.figure=self.fig\n        self.fig.axes.append(ax)\n        self.fig.add_axes(ax)\n        ax._subplotspec = gs\n        ax.set_position(gs.get_position(self.fig))\n        ax.set_subplotspec(gs)\n\n    def _finalize(self):\n        plt.close(self.sg.fig)\n        self.fig.canvas.mpl_connect(\"resize_event\", self._resize)\n        self.fig.canvas.draw()\n\n    def _resize(self, evt=None):\n        self.sg.fig.set_size_inches(self.fig.get_size_inches())","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-09-04T09:10:43.284023Z","iopub.execute_input":"2023-09-04T09:10:43.284443Z","iopub.status.idle":"2023-09-04T09:10:43.298772Z","shell.execute_reply.started":"2023-09-04T09:10:43.284413Z","shell.execute_reply":"2023-09-04T09:10:43.297726Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def _get_coll_root(coll: str) -> Path:\n    \"\"\"Parse the collection and return the corresponding data root.\n    \n    Parameters:\n        coll: collection\n\n    Return\n        data_root: data root of the collection\n    \"\"\"\n    coll_terms = coll.split(\"-\")\n    if len(coll_terms) == 3:\n        optim, src, search = coll_terms\n        data_root = DATA_ROOT/f\"{optim}/{src}/{search}\"\n    else:\n        optim, src = coll_terms\n        data_root = DATA_ROOT/f\"{optim}/{src}\"\n\n    return data_root\n\n\ndef _summarize(\n    df: pd.DataFrame,\n    file_name: Optional[str] = None,\n    n_rows_to_display: Optional[int] = 5,\n) -> None:\n    \"\"\"Summarize DataFrame.\n\n    Parameters:\n        df: input data\n        file_name: name of the input file\n        n_rows_to_display: number of rows to display\n\n    Return:\n        None\n    \"\"\"\n    file_name = \"Data\" if file_name is None else file_name\n\n    # Derive NaN ratio for each column\n    nan_ratio = pd.isna(df).sum() / len(df) * 100\n    nan_ratio.sort_values(ascending=False, inplace=True)\n    nan_ratio = nan_ratio.to_frame(name=\"NaN Ratio\").T\n\n    # Derive zero ratio for each column\n    zero_ratio = (df == 0).sum() / len(df) * 100\n    zero_ratio.sort_values(ascending=False, inplace=True)\n    zero_ratio = zero_ratio.to_frame(name=\"Zero Ratio\").T\n\n    # Print out summarized information\n    print(f\"=====Summary of {file_name}=====\")\n    display(df.head(n_rows_to_display))\n    print(f\"Shape: {df.shape}\")\n    print(\"NaN ratio:\")\n    display(nan_ratio)\n    print(\"Zero ratio:\")\n    display(zero_ratio)\n\n\ndef _plot_univar_dist(\n    data: Union[pd.Series, np.ndarray], feature: str, bins: int = 250, ax: Optional[Axes] = None\n) -> None:\n    \"\"\"Plot univariate distribution.\n\n    Parameters:\n        data: univariate data to plot\n        feature: feature name of the data\n        bins: number of bins\n        ax: user-specified axes\n\n    Return:\n        None\n    \"\"\"\n    if isinstance(data, np.ndarray):\n        data = pd.Series(data)\n\n    if ax is None:\n        fig, ax = plt.subplots(figsize=(8, 4))\n    sns.histplot(data=data, bins=bins, kde=True, palette=colors, ax=ax)\n    ax.axvline(x=data.mean(), color=\"orange\", linestyle=\"dotted\", linewidth=1.5, label=\"Mean\")\n    ax.axvline(\n        x=data.median(),\n        color=\"green\",\n        linestyle=\"dotted\",\n        linewidth=1.5,\n        label=\"Median\",\n    )\n    ax.axvline(\n        x=data.mode().values[0],\n        color=\"red\",\n        linestyle=\"dotted\",\n        linewidth=1.5,\n        label=\"Mode\",\n    )\n    ax.set_title(\n        f\"{feature.upper()} Distibution\\n\"\n        f\"Min {round(data.min(), 2)} | \"\n        f\"Max {round(data.max(), 2)} | \"\n        f\"Skewness {round(data.skew(), 2)} | \"\n        f\"Kurtosis {round(data.kurtosis(), 2)}\"\n    )\n    ax.set_xlabel(f\"{feature}\")\n    ax.set_ylabel(\"Bin Count\")\n    ax.legend()\n    if ax is None:\n        plt.show()\n        \n\ndef _plot_bivar(\n    data: Union[pd.Series, np.ndarray],\n    features: Optional[List[str]] = [\"0\", \"1\"],\n) -> Tuple[JointGrid, str]:\n    \"\"\"Plot bivariate distribution with regression line fitted.\n\n    Parameters:\n        data: bivariate data to plot\n        features: list of feature names\n\n    Return:\n        jg: seaborn JointGrid\n        title: figure title\n    \"\"\"\n    if isinstance(data, np.ndarray):\n        data = pd.DataFrame(data)\n    f1, f2 = features[0], features[1]\n    corr = data[[f1, f2]].corr().iloc[0, 1]\n\n    title = f\"{f1} versus {f2}, Corr={corr:.2}\"\n    jg = sns.jointplot(\n        x=data[f1],\n        y=data[f2],\n        kind=\"reg\",\n        height=6,\n        marginal_ticks=True,\n        joint_kws={\"line_kws\": {\"color\": \"orange\"}},\n    )\n    jg.ax_joint.set_xlabel(f1)\n    jg.ax_joint.set_ylabel(f2)\n    \n    return jg, title","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-09-04T09:10:46.042830Z","iopub.execute_input":"2023-09-04T09:10:46.043333Z","iopub.status.idle":"2023-09-04T09:10:46.058394Z","shell.execute_reply.started":"2023-09-04T09:10:46.043288Z","shell.execute_reply":"2023-09-04T09:10:46.056745Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"data_appearance\"></a>\n## 1. Data Appearance\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\nAt first, I find it hard to understand how the raw data is structured. Hence, I read the amazing [paper](https://arxiv.org/abs/2308.13490) provided by the host, which clearly explains how the data is generated and arranged.\n\nIn short, the hierarchy is defined as *optimization-source-search-split*. That is, collections of data differ in terms of:\n1. *optimization* - The compiler optimization\n    * `layout`: The tensor computational graph is the input graph to the **layoyt assignment pass**.\n    * `tile`: The fused subgraph (**kernel**) will goes through **tile size selection**.\n2. *source* - The source of graphs (graph collection)<br>\nThere are two sources from which the computational graphs are collected:\n    * `xla`: The combination of the XLA regression benchmark.\n    * `nlp`: A variety of BERT for training and inference.\n3. *search* - The search strategy (configuration generation)<br>\nThe XLA autotuner is used to generate the configurations, and there are two strategies to explore the search space:\n    * `default`: Explore the search space using a **genetic algorithm** starting from the default configuration.\n    * `random`: Pick random candidates in the search space.\n4. *split* - The dataset split (*i.e.,* `train`, `valid`, and `test`)\n  \nWith the information above, let's take a look at the data folder structure. Some observations are summarized as follows:\n* `tile` optimization contains graphs only from `xla` source.\n* There's only one search strategy for `tile` optimization (at the level of the **fused subgraph**).\n    * **Graph-level** optimization is run beforehand to output a collection of kernels.\n    * Configurations are generated by **enumerating** all possible tile sizes for the kernel in a random order.    ","metadata":{}},{"cell_type":"code","source":"!tree -d /kaggle/input/predict-ai-model-runtime/npz_all/npz","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:10:47.915415Z","iopub.execute_input":"2023-09-04T09:10:47.915839Z","iopub.status.idle":"2023-09-04T09:10:50.781302Z","shell.execute_reply.started":"2023-09-04T09:10:47.915807Z","shell.execute_reply":"2023-09-04T09:10:50.780091Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"ml_optim\"></a>\n## 2. ML Compiler Optimizations\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\nThere are two types of optimizations commonly used in ML compilers, namely **graph-level** and **kernel-level** optimizations. Let's explore each of them with the provided data.\n\n<a id=\"graph_lv\"></a>\n### *Graph-Level Optimizations*\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\n[![Screenshot-2023-09-03-at-15-47-46-Google.png](https://i.postimg.cc/x1y5Bj0C/Screenshot-2023-09-03-at-15-47-46-Google.png)](https://postimg.cc/21ynZDLR)\n> Graph-level optimizations require the entire program graph to make decisions.\n\nAt this level, the input of the optimization processes is **the entire computational graph** and the output are a collection of **fused subgraphs (kernels)**. In this comptition, **layout assignment pass** is selected as the focus.\n\n> *Layout* configurations control how tensors are laid out in the physical memory, by specifying the dimension order of each input and output of an operation node.\n\nAfter peeking a `layout` example, we can see that each `.npz` file stores many fields to be explored further.","metadata":{}},{"cell_type":"code","source":"layout_demo = dict(np.load(DATA_ROOT/\"layout/nlp/default/train/albert_en_base_batch_size_16_test.npz\"))\nfor k, v in layout_demo.items():\n    print(f\"{k:<16}\", f\"| Type: {type(v)} | Dtype: {str(v.dtype):<8} | Shape {v.shape}\")","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:10:51.119185Z","iopub.execute_input":"2023-09-04T09:10:51.119622Z","iopub.status.idle":"2023-09-04T09:10:52.904575Z","shell.execute_reply.started":"2023-09-04T09:10:51.119589Z","shell.execute_reply":"2023-09-04T09:10:52.903469Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Before stepping further, let's define some commonly used notations as follows:\n* `n`: number of nodes (*i.e.,* tensor operations) in the graph.\n* `m`: number of edges (*i.e.,* tensor flows) in the graph.\n* `c`: number of differenct configurations for the specific graph.\n* `nc`: number of **configurable** nodes in the graph, must be less than or equal to `n`.\n\n####  a. `node_feat`\n`node_feat` is an `(n, 140)` feature array, in which the nodes are ordered **topologically**. Each column has its own meaning, which can be found in the [data tab](https://www.kaggle.com/competitions/predict-ai-model-runtime/data). The source code used to extract feature can be found [here](https://github.com/google-research-datasets/tpu_graphs/blob/main/tpu_graphs/process_data/xla/featurizers.h#L542).","metadata":{}},{"cell_type":"code","source":"n = layout_demo[\"node_feat\"].shape[0]\nprint(f\"There are {n} nodes in this demo graph.\")","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:10:53.069949Z","iopub.execute_input":"2023-09-04T09:10:53.070366Z","iopub.status.idle":"2023-09-04T09:10:53.077485Z","shell.execute_reply.started":"2023-09-04T09:10:53.070307Z","shell.execute_reply":"2023-09-04T09:10:53.075731Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Before proceeding to the next field, let's do some basic data checkings, including:\n1. Whether column 1 of `node_feat` is deprecated for all files.\n    * As described in the **Node Features** section of [data tab](https://www.kaggle.com/competitions/predict-ai-model-runtime/data).\n2. Whether each node has its own op-code.\n    * Because each node represents a tensor operation, one node should map to one of the op-codes.\n3. Whether the number of configurable nodes is always **less than or equal to** the total number of nodes.\n    * As configurable node set is the subset of the entire node set, we can expect `nc <= n`.","metadata":{}},{"cell_type":"code","source":"data_checker_layout = DataChecker(\"layout\")\nnpz_illeg_layout = data_checker_layout.check()","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:11:02.957759Z","iopub.execute_input":"2023-09-04T09:11:02.958183Z","iopub.status.idle":"2023-09-04T09:25:07.245078Z","shell.execute_reply.started":"2023-09-04T09:11:02.958147Z","shell.execute_reply":"2023-09-04T09:25:07.242497Z"},"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Because column 1, `element_size_in_bits`, is deprecated for all layouts, we can safely drop this column.","metadata":{}},{"cell_type":"code","source":"layout_with_nonzero_col1 = npz_illeg_layout[\"col1_deprecated\"]\nprint(f\"There are {len(layout_with_nonzero_col1)} layouts whose column 1 of node_feat has non-zeros.\")","metadata":{"_kg_hide-input":false,"execution":{"iopub.status.busy":"2023-09-04T09:25:07.249627Z","iopub.execute_input":"2023-09-04T09:25:07.250074Z","iopub.status.idle":"2023-09-04T09:25:07.260794Z","shell.execute_reply.started":"2023-09-04T09:25:07.250034Z","shell.execute_reply":"2023-09-04T09:25:07.258767Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### b. `node_opcode`\n`node_opcode` is an `(n, )` array, whose entries represent **op-codes** of the corresponding nodes. The mapping of op-codes and instructions can be found [here](https://github.com/google-research-datasets/tpu_graphs/blob/main/tpu_graphs/process_data/xla/hlo_opcode.h#L94).\n\nAll layouts have perfect op-code mapping.","metadata":{}},{"cell_type":"code","source":"layout_with_unaligned_ops = npz_illeg_layout[\"op_align\"]\nprint(f\"There are {len(layout_with_unaligned_ops)} layouts with unaligned op-codes.\")","metadata":{"_kg_hide-input":false,"execution":{"iopub.status.busy":"2023-09-04T09:25:07.263518Z","iopub.execute_input":"2023-09-04T09:25:07.264165Z","iopub.status.idle":"2023-09-04T09:25:07.284324Z","shell.execute_reply.started":"2023-09-04T09:25:07.264113Z","shell.execute_reply":"2023-09-04T09:25:07.282697Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### c. `edge_index`\n`edge_index` is an `(m, 2)` array, each row of which represents a directed edge `[u, v]` pointing from node `u` to `v`.\n\n<div class=\"alert alert-block alert-warning\">\n    <h5>As described in the <a href=\"https://www.kaggle.com/competitions/predict-ai-model-runtime/data\">data tab</a>, point <code>u</code> consumes the output of <code>v</code>. But, I think the opposite is true, isn't it?</h5>\n</div>\n<div class=\"alert alert-block alert-danger\">\n    <h5>The competition host has clarified the issue in <a href=\"https://www.kaggle.com/competitions/predict-ai-model-runtime/discussion/436786\">this forum</a>. Also, the description about <code>edge_index</code> in the data tab has been modified.</h5>\n</div>","metadata":{}},{"cell_type":"markdown","source":"#### d. `node_config_ids`\n> The layout configuration of a graph is a collection of per-node layout decisions on configurable\nnodes (*i.e.,* convolution and reshape).\n\n`node_config_ids` is an `(nc, )` array, whose entries represent the indices of **configurable** nodes.\n\nAll layouts have reasonable number of configurable nodes.","metadata":{}},{"cell_type":"code","source":"layout_with_more_cfg_nodes = npz_illeg_layout[\"n_cfg_nodes\"]\nprint(f\"There are {len(layout_with_more_cfg_nodes)} layouts with more configurable nodes than total.\")","metadata":{"_kg_hide-input":false,"execution":{"iopub.status.busy":"2023-09-04T09:25:07.288619Z","iopub.execute_input":"2023-09-04T09:25:07.289255Z","iopub.status.idle":"2023-09-04T09:25:07.303441Z","shell.execute_reply.started":"2023-09-04T09:25:07.289215Z","shell.execute_reply":"2023-09-04T09:25:07.302360Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### e. `node_config_feat`\n`node_config_feat` is an `(c, nc, 18)` configuration feature array. For an entry `[j, k]`, it represents the configuration feature vector of the configurable node `layout_tmp[\"node_config_ids\"][k]` for the `j`-th compilation run. Each feature field has its own meaning, which can be found in the [data tab](https://www.kaggle.com/competitions/predict-ai-model-runtime/data).\n\nFollowing is a quick illustration of the relationship between `config`-related fields. Let's take `j == 0` and `k == 2`.","metadata":{}},{"cell_type":"code","source":"c = layout_demo[\"node_config_feat\"].shape[0]\nj, k = 0, 2\nnode_config_id = layout_demo[\"node_config_ids\"][k]\nnode_config_feat_demo = layout_demo[\"node_config_feat\"][j][k]\n\nprint(f\"There are {c} different configurations for this graph.\")\nprint(f\"The node configuration feature of the configurable node {node_config_id} in the {j}-th compilation run is:\")\nprint(node_config_feat_demo)","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:25:07.304701Z","iopub.execute_input":"2023-09-04T09:25:07.305694Z","iopub.status.idle":"2023-09-04T09:25:07.320250Z","shell.execute_reply.started":"2023-09-04T09:25:07.305619Z","shell.execute_reply":"2023-09-04T09:25:07.318754Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### f. `config_runtime`\n`config_runtime` is an `(c, )` target array, where the `j`-th entry represents the runtime of the `j`-th compilation run.\n\nThis field can be seen as the target of the regression task if you want to predict the runtime directly. Following shows the distribution of `config_runtime` of the demo example. As can be observed, there seems to exist multiple modals in different regions in `default` search strategy. As stated in the paper,\n\n>  The default collection tends to contain configurations that are not too different from the default, and have similar execution times.\n\nAs for `random`, the runtimes seem to spread out, and don't show the clustering effect.\n\n>  The random collection includes very different configurations with very different execution times.","metadata":{}},{"cell_type":"code","source":"layout_demo2 = dict(np.load(DATA_ROOT/\"layout/nlp/random/train/albert_en_base_batch_size_16_test.npz\"))\n\nfig, axes = plt.subplots(nrows=1, ncols=2, figsize=(16, 4))\n_plot_univar_dist(layout_demo[\"config_runtime\"], \"Runtimes of Diff Configs (layout - default)\", 1000, ax=axes[0])\n_plot_univar_dist(layout_demo2[\"config_runtime\"], \"Runtimes of Diff Configs (layout - random)\", 1000, ax=axes[1])\nplt.tight_layout()","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:25:07.322621Z","iopub.execute_input":"2023-09-04T09:25:07.323745Z","iopub.status.idle":"2023-09-04T09:25:16.804127Z","shell.execute_reply.started":"2023-09-04T09:25:07.323696Z","shell.execute_reply":"2023-09-04T09:25:16.802607Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### g. `node_splits`\n`node_splits` provides the information about the starting of HLO computations in the graph, which is beneficial to graph segmentations. Nodes `layout_tmp[\"node_splits\"][i]` to `layout_tmp[\"node_splits\"][i+1] - 1` belongs to the same computation.\n\nIn the demo examples, nodes `0` to `1095` belong to the same computation. That is, the entire graph belongs to the same computation.","metadata":{}},{"cell_type":"code","source":"print(f\"The splitting point of the demo layout is {layout_demo['node_splits']}.\")","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:25:16.806610Z","iopub.execute_input":"2023-09-04T09:25:16.807620Z","iopub.status.idle":"2023-09-04T09:25:16.814939Z","shell.execute_reply.started":"2023-09-04T09:25:16.807573Z","shell.execute_reply":"2023-09-04T09:25:16.813761Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"kernel_lv\"></a>\n### *Kernel-Level Optimizations*\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\n[![Screenshot-2023-09-03-at-19-12-38-Google.png](https://i.postimg.cc/9QDZDp35/Screenshot-2023-09-03-at-19-12-38-Google.png)](https://postimg.cc/mPGtq728)\n> Kernel-level optimizations can be done independently within each kernel.\n\nAt this level, a **kernel** (represented as a **fused subgraph**) is a fusion of multiple tensor oprations (*i.e.,* nodes in the computational graph). Kernels are optimized one at a time, independent of all the others. In this competition, **tile size selection** is our target compiler optimization.\n\n>  *Tile* configuration controls the tile size of each fused subgraph\n\nJust as `layout`, there are also many fields to be explored in `tile`. Furthermore, there exist many common fields. However, the `config`-related fields look a little bit different. Let's see what's going on.","metadata":{}},{"cell_type":"code","source":"tile_demo = dict(np.load(DATA_ROOT/\"tile/xla/train/alexnet_train_batch_32_-1bae27a41d70f4dc.npz\"))\nfor k, v in tile_demo.items():\n    print(f\"{k:<27}\", f\"| Type: {type(v)} | Dtype: {str(v.dtype):<8} | Shape {v.shape}\")","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:25:16.816348Z","iopub.execute_input":"2023-09-04T09:25:16.816710Z","iopub.status.idle":"2023-09-04T09:25:16.846382Z","shell.execute_reply.started":"2023-09-04T09:25:16.816673Z","shell.execute_reply":"2023-09-04T09:25:16.845554Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"First, the quick checkings show that `tile` also follows the statement in the [data tab](https://www.kaggle.com/competitions/predict-ai-model-runtime/data); that is, the column 1 of `node_feat` is deprecated. Also, the alignment of op-codes is fine.","metadata":{}},{"cell_type":"code","source":"data_checker_tile = DataChecker(\"tile\")\nnpz_illeg_tile = data_checker_tile.check()","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-09-04T09:25:16.847704Z","iopub.execute_input":"2023-09-04T09:25:16.848205Z","iopub.status.idle":"2023-09-04T09:26:21.806063Z","shell.execute_reply.started":"2023-09-04T09:25:16.848175Z","shell.execute_reply":"2023-09-04T09:26:21.804772Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tile_with_nonzero_col1 = npz_illeg_tile[\"col1_deprecated\"]\ntile_with_unaligned_ops = npz_illeg_tile[\"op_align\"]\nprint(f\"There are {len(tile_with_nonzero_col1)} tiles whose column 1 of node_feat has non-zeros.\")\nprint(f\"There are {len(tile_with_unaligned_ops)} tiles with unaligned op-codes.\")","metadata":{"_kg_hide-output":false,"execution":{"iopub.status.busy":"2023-09-04T09:26:21.809332Z","iopub.execute_input":"2023-09-04T09:26:21.809713Z","iopub.status.idle":"2023-09-04T09:26:21.817003Z","shell.execute_reply.started":"2023-09-04T09:26:21.809680Z","shell.execute_reply":"2023-09-04T09:26:21.815751Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"As state in the [data tab](https://www.kaggle.com/competitions/predict-ai-model-runtime/data),\n\n> Crucially, the configuration is at the graph-level.\n\nWe already know that **tile size selection** is done as **kernel-level optimization**. Hence, the description actually means that the `tile` configuration is a configuration for the **entired fused subgraph**, which is different from the **per-node** manner in **graph-level optimization**.","metadata":{}},{"cell_type":"markdown","source":"#### a. `config_feat`\n`config_feat` is a `(c, 24)` array, each row of which is the **(sub)graph-level** configuration feature vector. Each feature field has its own meaning, which can be found in the [data tab](https://www.kaggle.com/competitions/predict-ai-model-runtime/data).\n\nFollowing is a quick illustration of the `j`-th configuration feature vector, where we take `j == 0`.","metadata":{}},{"cell_type":"code","source":"c = tile_demo[\"node_feat\"].shape[0]\nj = 0\nconfig_feat_demo = tile_demo[\"node_feat\"][j]\n\nprint(f\"There are {c} different configurations for this fused subgraph (kernel).\")\nprint(f\"The {j}-th configuration feature vector is:\")\nprint(config_feat_demo.astype(np.int32))","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:26:28.437584Z","iopub.execute_input":"2023-09-04T09:26:28.438007Z","iopub.status.idle":"2023-09-04T09:26:28.446361Z","shell.execute_reply.started":"2023-09-04T09:26:28.437976Z","shell.execute_reply":"2023-09-04T09:26:28.445078Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### b. `config_runtime` & `config_runtime_normalizers`\n`config_runtime` and `config_runtime_normalizers` are both `(c, )` target arrays, where the former represents the runtime of the graphs compiled with the corresponding configurations and the latter with the **default** configuration. Slightly different from `layout`, we need to retrieve the configurations with the smallest `tile_tmp[\"config_runtime\"] / tile_tmp[\"config_runtime_normalizers\"]`, indicating the real runtime compared with the default setting (in the denominator). \n\nFollowing shows the distribution of `config_runtime` and runtime ratio of the demo example. Though they look the same in this example, it's possible to observe the difference between the two. And, the reason behind the scene is stated in the [data tab](https://www.kaggle.com/competitions/predict-ai-model-runtime/data),\n\n> Samples from the same graph may have slightly different `\"config_runtime_normalizers\"` because they are measured from different runs on multiple machines.","metadata":{}},{"cell_type":"code","source":"ratio = tile_demo[\"config_runtime\"] / tile_demo[\"config_runtime_normalizers\"]\n\nfig, axes = plt.subplots(nrows=1, ncols=2, figsize=(16, 4))\n_plot_univar_dist(tile_demo[\"config_runtime\"], \"Runtimes of Diff Configs (tile)\", 50, ax=axes[0])\n_plot_univar_dist(ratio, \"Runtime Ratios of Diff Configs (tile)\", 50, ax=axes[1])\nplt.tight_layout()","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:26:31.594286Z","iopub.execute_input":"2023-09-04T09:26:31.594717Z","iopub.status.idle":"2023-09-04T09:26:32.922306Z","shell.execute_reply.started":"2023-09-04T09:26:31.594681Z","shell.execute_reply":"2023-09-04T09:26:32.921070Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"graph_stats\"></a>\n## 3. Basic Graph Statistics\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\nTo better understand the properties of computational graphs in hand, let's derive basic graph statistics in a **collection-specific** manner. Then, we can compare if there exists different characteristic among collections. More stats can be derived if you're interested.\n\nCurrently implemented statistics are listed as follows:\n1. `n_nodes`: Number of nodes in the computational graph, indicating how many **tensor oprations** are performed.\n2. `n_edges`: Number of edges in the computational graph, indicating the total number of **tensor flows**.\n3. `avg_deg`: Average degree of the computational graph, which is defined as,\n$$\n    avg\\_deg = \\frac{1}{n} (\\sum_{i}^{n} in\\_deg_i + \\sum_{i}^{n} out\\_deg_i)\n$$\n, where $n$ is the total number of nodes, $in\\_deg_i$ and $out\\_deg_i$ are in-degree and out-degree of node $i$.\n4. `avg_clust_coeff`: Average clustering coefficient of the computational graph, which is defined as a proportion of the number of **directed triangles** through the anchor node $u$ divided by the number of **possible links** within $u$'s neighbourhood.\n    * For implementation, please see the source code [here](https://networkx.org/documentation/stable/_modules/networkx/algorithms/cluster.html#average_clustering).\n5. `longest_path`: The longest path in the computational graph, indicating how deep the model architecture is.\n\nFollowing illustrates what is the **directed triangles**. For each node acting as an anchor (take `conv` as an example), 2 directed triangles can be formed from the perspective of its two neighbors (*e.g.,* `reshape` and `add`). Finding different directed polygons (*e.g.,* rectangle) can somewhat help us better understand the model architecture (*e.g.,*  residual links).","metadata":{}},{"cell_type":"code","source":"edge_list = [(0, 1), (1, 2), (0, 2)]\ng = nx.DiGraph(edge_list)\n\nfig, ax = plt.subplots(figsize=(8, 4))\noptions = {\n    \"node_color\": \"tab:blue\",\n    'node_size': 2500,\n    \"node_shape\": \"s\",\n    \"width\": 3,\n    \"arrowstyle\": \"-|>\",\n    \"arrowsize\": 12,\n}\npos = {0: [0.2, 0.2], 1: [0.5, 0.5], 2: [0.8, 0.2]}\nlabels = {0: \"reshape\", 1: \"conv\", 2: \"add\"}\nnx.draw_networkx(g, with_labels=False, pos=pos, ax=ax, arrows=True, **options)\nnx.draw_networkx_labels(g, pos, labels, ax=ax, font_size=12, font_color=\"whitesmoke\")\nplt.tight_layout()\nplt.axis(\"off\")\nplt.show()","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class GraphStatsEngineer(object):\n    \"\"\"Derive basic statistics of graphs from different collections.\n\n    Parameters:\n        coll: collection\n            *Note: Run all collections if not given\n        debug: if False, sanity checks are ignored\n    \"\"\"\n\n    def __init__(self, coll: Optional[str] = None, debug: bool = False) -> None:\n        self.coll = [coll] if coll is not None else COLL\n        self.debug = debug\n\n    def run(self, stats_list: List[str] = [\"n_nodes\"]) -> Tuple[Dict[str, Any], List[str]]:\n        \"\"\"Derive the graph stats.\n        \n        Parameters:\n            stats_list: graph stats, the choices are as follows:\n                {\"n_nodes\", \"n_edges\", \"avg_deg\", \"avg_clust_coeff\", \"longest_path\"}\n\n        Return:\n            stats: graph statistics\n            npz_illeg: .npz files failing sanity checks\n        \"\"\"\n        stats = {}\n        npz_illeg = []\n        \n        for coll in self.coll:\n            stats[coll] = {\"split\": [], \"file\": [], \"n_graphs\": 0}\n            for s in stats_list:\n                stats[coll][s] = []\n            coll_root = _get_coll_root(coll)\n            \n            for split in SPLIT:\n                data_root = coll_root/split\n                for file in tqdm(os.listdir(data_root)):\n                    if file.endswith(\".npz\"):\n                        stats[coll][\"split\"].append(split)\n                        stats[coll][\"file\"].append(file)\n                        stats[coll][\"n_graphs\"] += 1\n                    else:\n                        continue\n                    \n                    # Prepare data\n                    npz_tmp = self._load_npz(data_root/file)\n                    g_tmp = self._get_digraph(npz_tmp[\"edge_index\"])\n    \n                    # Derive stats\n                    stats_tmp = self._run_single(npz_tmp, g_tmp, stats_list)\n                    if stats_tmp is not None:\n                        for s, v in stats_tmp.items():\n                            stats[coll][s].append(v)\n                    else:\n                        npz_illeg.append(str(data_root/file))\n                        \n        return stats, npz_illeg\n\n    def _load_npz(self, data_file: Path) -> Dict[str, np.ndarray]:\n        \"\"\"Load and return the graph with its configurations.\n\n        Parameters:\n            data_file: path to data\n\n        Return:\n            npz_tmp: graph and configurations\n        \"\"\"\n        npz_tmp = dict(np.load(data_file))\n\n        return npz_tmp\n    \n    def _get_digraph(self, edge_index: np.ndarray) -> nx.DiGraph:\n        \"\"\"Return the NetworkX Graph.\n\n        Parameters:\n            edge_index: edge index\n\n        Return:\n            digraph: directed graph representation of the computational\n                graph\n        \"\"\"\n        edge_list = list(map(tuple, edge_index))\n        digraph = nx.DiGraph(edge_list)\n\n        return digraph\n    \n    def _run_single(\n        self,\n        npz: Dict[str, np.ndarray],\n        g: nx.DiGraph,\n        stats_list: List[str]\n    ) -> Optional[Dict[str, float]]:\n        \"\"\"Derive the graph stats for a single graph.\n        \n        Parameters:\n            npz: graph and configurations\n            g: directed graph representation of the computational graph\n            stats: graph stats\n\n        Return:\n            stats_single: graph stats of the current graph\n        \"\"\"\n        stats_single = {}\n        if \"n_nodes\" in stats_list:\n            n_nodes = g.order()\n            if self.debug and n_nodes != len(npz[\"node_feat\"]):\n                return None\n            stats_single[\"n_nodes\"] = n_nodes\n        if \"n_edges\" in stats_list:\n            n_edges = g.size()\n            if self.debug and n_edges != len(npz[\"edge_index\"]):\n                return None\n            stats_single[\"n_edges\"] = n_edges\n        if \"avg_deg\" in stats_list:\n            n_nodes = g.order()  # Note if any orphan node\n            avg_deg = sum(dict(g.degree).values()) / n_nodes\n            avg_indeg = sum(dict(g.in_degree).values()) / n_nodes\n            avg_outdeg = sum(dict(g.out_degree).values()) / n_nodes\n            if self.debug and avg_deg != avg_indeg + avg_outdeg:\n                return None\n            stats_single[\"avg_deg\"] = avg_deg\n        if \"avg_clust_coeff\" in stats_list:\n            avg_clust_coeff = nx.average_clustering(g)\n            stats_single[\"avg_clust_coeff\"] = avg_clust_coeff\n        if \"longest_path\" in stats_list:\n            longest_path = nx.dag_longest_path_length(g)\n            stats_single[\"longest_path\"] = longest_path\n\n        return stats_single","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-09-04T09:26:50.535995Z","iopub.execute_input":"2023-09-04T09:26:50.536452Z","iopub.status.idle":"2023-09-04T09:26:50.557245Z","shell.execute_reply.started":"2023-09-04T09:26:50.536417Z","shell.execute_reply":"2023-09-04T09:26:50.555494Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"eng = GraphStatsEngineer(debug=False)\nstats_list = [\"n_nodes\", \"n_edges\", \"avg_deg\", \"avg_clust_coeff\", \"longest_path\"]\nstats, _ = eng.run(stats_list)","metadata":{"execution":{"iopub.status.busy":"2023-09-04T09:26:58.522348Z","iopub.execute_input":"2023-09-04T09:26:58.522829Z","iopub.status.idle":"2023-09-04T09:28:30.986481Z","shell.execute_reply.started":"2023-09-04T09:26:58.522796Z","shell.execute_reply":"2023-09-04T09:28:30.984978Z"},"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"After comparing graph statistics in different collections, some observations are summarized as follows:\n1. For optimization `layout`, distributions of stats `avg_deg` and `avg_clust_coeff` of source `xla` are more different for different search strategies than source `nlp`.\n2. There exist much deeper model architectures in `layout-xla` compared with `layout-nlp`.\n3. `tile-xla` shows the most distinct characteristics.\n    * It has the most number of graphs compared with the others.\n    * Graphs have much less number of nodes and edges.\n    * Longest paths are much shorter.\n\nPoint 3 is reasonable, because graphs in `tile-xla` are **fused subgraphs**, which are intrinsically smaller than the entire graph.","metadata":{}},{"cell_type":"markdown","source":"<a id=\"lnd\"></a>\n### *`layout-nlp-default`*\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)","metadata":{}},{"cell_type":"code","source":"coll = \"layout-nlp-default\"\n\nprint(f\"There are {stats[coll]['n_graphs']} graphs in collection {coll}.\")\nfig, axes = plt.subplots(nrows=1, ncols=5, figsize=(35, 4))\nfor i, s in enumerate(stats_list):\n    _plot_univar_dist(np.array(stats[coll][s]), f\"{coll} - {s}\", 50, ax=axes[i])\nplt.tight_layout()","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"lnr\"></a>\n### *`layout-nlp-random`*\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)","metadata":{}},{"cell_type":"code","source":"coll = \"layout-nlp-random\"\n\nprint(f\"There are {stats[coll]['n_graphs']} graphs in collection {coll}.\")\nfig, axes = plt.subplots(nrows=1, ncols=5, figsize=(35, 4))\nfor i, s in enumerate(stats_list):\n    _plot_univar_dist(np.array(stats[coll][s]), f\"{coll} - {s}\", 50, ax=axes[i])\nplt.tight_layout()","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"lxd\"></a>\n### *`layout-xla-default`*\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)","metadata":{}},{"cell_type":"code","source":"coll = \"layout-xla-default\"\n\nprint(f\"There are {stats[coll]['n_graphs']} graphs in collection {coll}.\")\nfig, axes = plt.subplots(nrows=1, ncols=5, figsize=(35, 4))\nfor i, s in enumerate(stats_list):\n    _plot_univar_dist(np.array(stats[coll][s]), f\"{coll} - {s}\", 50, ax=axes[i])\nplt.tight_layout()","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"lxr\"></a>\n### *`layout-xla-random`*\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)","metadata":{}},{"cell_type":"code","source":"coll = \"layout-xla-random\"\n\nprint(f\"There are {stats[coll]['n_graphs']} graphs in collection {coll}.\")\nfig, axes = plt.subplots(nrows=1, ncols=5, figsize=(35, 4))\nfor i, s in enumerate(stats_list):\n    _plot_univar_dist(np.array(stats[coll][s]), f\"{coll} - {s}\", 50, ax=axes[i])\nplt.tight_layout()","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"tx\"></a>\n### *`tile-xla`*\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)","metadata":{}},{"cell_type":"code","source":"coll = \"tile-xla\"\n\nprint(f\"There are {stats[coll]['n_graphs']} graphs in collection {coll}.\")\nfig, axes = plt.subplots(nrows=1, ncols=5, figsize=(35, 4))\nfor i, s in enumerate(stats_list):\n    _plot_univar_dist(np.array(stats[coll][s]), f\"{coll} - {s}\", 250, ax=axes[i])\nplt.tight_layout()","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"There exist some files failing the sanity checks in stats engineering process (please set `debug` to `True` if you want to explore further). After analysis, they all fail the check of `n_nodes == len(npz[\"node_feat\"])`. The reason behind the scene is that I use `edge_index` to construct the `DiGraph`. Hence, isolated nodes aren't taken into consideration.\n\n<div class=\"alert alert-block alert-warning\">\n    <h5>But, what does an isolated node mean in the computational graph?</h5>\n</div>\n<div class=\"alert alert-block alert-danger\">\n    <h5>The meaning of isolated nodes in a computational graph has been explained by the host <a href=\"https://www.kaggle.com/competitions/predict-ai-model-runtime/discussion/437068#2425354\">here</a>. Also, I've confirmed that all isolated nodes in the raw data are <b>parameter</b> (with op-code 63).</h5>\n</div>","metadata":{}},{"cell_type":"code","source":"# for data_file in npz_illeg:\n#     print(f\"File {data_file}\")\n#     npz_tmp = dict(np.load(data_file))\n#     g_tmp = nx.DiGraph(list(map(tuple, npz_tmp[\"edge_index\"])))\n\n#     iso_nodes = set(np.arange(len(npz_tmp[\"node_feat\"]))).difference(set(g_tmp.nodes()))\n#     print(f\"-> Isolated nodes: {iso_nodes}\")\n\n#     # Check op-codes of isolated nodes\n#     iso_op_codes = npz_tmp[\"node_opcode\"][list(iso_nodes)]\n#     iso_op_codes_uniq = np.unique(iso_op_codes)\n#     assert len(iso_op_codes_uniq) == 1 and iso_op_codes_uniq[0] == 63","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Let's convert graph stats to `DataFrame` and dump it `.csv` file for quicker access in the future.","metadata":{}},{"cell_type":"code","source":"df = []\nfor coll, stats_ in stats.items():\n    df_coll = pd.DataFrame.from_dict(stats_)\n    df_coll[\"coll\"] = coll\n    df.append(df_coll)\ndf = pd.concat(df, ignore_index=True)\ndf = df[[df.columns[-1]] + [c for c in df.columns[:-1]]]\n_summarize(df, \"Graph Stats\", 2)\ndf.to_csv(\"graph_stats.csv\", index=False)","metadata":{"_kg_hide-input":false},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"runtime\"></a>\n## 4. Our Target - Runtime\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\nIn this competition, our goal is to predict the runtime of graphs and configurations (*i.e.,* execution time of the graph compiled with the specific configuration). Hence, let's move on to the main focus, `config_time`, to get familiar with it! ","metadata":{}},{"cell_type":"code","source":"fig, axes = plt.subplots(nrows=1, ncols=len(COLL), figsize=(35, 4))\nfor i, coll in enumerate(COLL):\n    coll_root_tr = _get_coll_root(coll)/\"train\"\n    data_files = os.listdir(coll_root_tr)\n    data_file = coll_root_tr/data_files[random.randint(0, len(data_files))]\n    npz_tmp = dict(np.load(data_file))\n\n    n_configs = len(npz_tmp[\"config_runtime\"])\n    if coll != \"tile-xla\":\n        runtime = npz_tmp[\"config_runtime\"]\n    else:\n        runtime = npz_tmp[\"config_runtime\"] / npz_tmp[\"config_runtime_normalizers\"]\n    _plot_univar_dist(runtime, f\"{coll} Runtime Demo (#Configs = {n_configs})\", ax=axes[i])\nplt.tight_layout()","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"To analyze runtime behavior, we first derive runtime stats with the way similar to graph stats derivation. Then, we'll combine graph stats with runtime stats to do some comparisons.","metadata":{}},{"cell_type":"code","source":"def _get_runtime_stats(runtime: np.ndarray) -> Dict[str, float]:\n    \"\"\"Derive and return runtime stats of one computational graph.\n\n    Parameters:\n        runtime: runtime of different configs\n\n    Return:\n        runtime_stats: runtime stats of one graph\n    \"\"\"\n    runtime_stats = {\n        \"runtime_mean\": np.mean(runtime),\n        \"runtime_std\": np.std(runtime),\n        \"runtime_min\": np.min(runtime),\n        \"runtime_max\": np.max(runtime),\n        \"runtime_median\": np.median(runtime),\n        \"runtime_mode\": mode(runtime),\n        \"runtime_skew\": skew(runtime),\n        \"runtime_kurt\": kurtosis(runtime)\n    }\n\n    return runtime_stats","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"runtime_stats_list = [\"mean\", \"std\", \"min\", \"max\",\n                      \"median\", \"mode\", \"skew\", \"kurt\"]\ntarget_stats = {}\nfor coll in COLL:\n    target_stats[coll] = {\"split\": [], \"file\": [], \"n_configs\": []}\n    for s in runtime_stats_list:\n        target_stats[coll][f\"runtime_{s}\"] = []\n    coll_root = _get_coll_root(coll)\n    \n    for split in SPLIT:\n        data_root = coll_root/split\n        for file in tqdm(os.listdir(data_root)):\n            if not file.endswith(\".npz\"):\n                continue\n\n            npz_tmp = dict(np.load(data_root/file))\n            target_stats[coll][\"split\"].append(split)\n            target_stats[coll][\"file\"].append(file)\n            target_stats[coll][\"n_configs\"].append(len(npz_tmp[\"config_runtime\"]))\n            \n            # Derive runtime stats\n            if coll != \"tile-xla\":\n                runtime = npz_tmp[\"config_runtime\"]\n            else:\n                runtime = npz_tmp[\"config_runtime\"] / npz_tmp[\"config_runtime_normalizers\"]\n            runtime_stats = _get_runtime_stats(runtime)\n            for s, v in runtime_stats.items():\n                target_stats[coll][s].append(v)","metadata":{"_kg_hide-output":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Again, we dump runtime stats `DataFrame` for future access.","metadata":{}},{"cell_type":"code","source":"df_target = []\nfor coll, target_stats_ in target_stats.items():\n    df_coll = pd.DataFrame.from_dict(target_stats_)\n    df_coll[\"coll\"] = coll\n    df_target.append(df_coll)\ndf_target = pd.concat(df_target, ignore_index=True)\ndf_target = df_target[[df_target.columns[-1]] + [c for c in df_target.columns[:-1]]]\n_summarize(df_target, \"Runtime Stats\", 2)\ndf_target.to_csv(\"runtime_stats.csv\", index=False)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Because we have no access to groundtruth of the test set, we keep only the training and validation sets for analysis. What's more, graphs in `layout` and `tile` are at different levels (*i.e.,* graph-level and kernel-level), having `config_runtime` at different scales. Hence, we'll analyze them separately.","metadata":{}},{"cell_type":"code","source":"df = df.merge(df_target, on=[\"coll\", \"split\", \"file\"])  # Merge graph stats with runtime stats\ndf = df[df[\"split\"] != \"test\"].reset_index(drop=True)\ndf_layout = df[df.coll.str.startswith(\"layout\")].reset_index(drop=True)\ndf_tile = df[df.coll.str.startswith(\"tile\")].reset_index(drop=True)\nprint(f\"There are {len(df_layout)} graphs in layout and {len(df_tile)} in tile.\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"To observe runtime behavior w.r.t. graph stats, we select `runtime_mean` (*i.e.,* average runtime over all configurations) as the target. Also, we take `np.log1p` on `runtime_mean` for better visualization and interpretation.\n\nIn general, computational graphs with the following properties tend to run longer:\n1. Have more nodes.\n    * More tensor operations to run.\n2. Have more edges.\n    * More tensor flows to pass.\n3. With greater average degrees.\n    * Probably more complex connetion in network.\n4. With longer path.\n    * Probabaly deeper network.","metadata":{}},{"cell_type":"code","source":"df_layout[\"runtime_mean_log\"] = np.log1p(df_layout.runtime_mean)\ndf_tile[\"runtime_mean_log\"] = np.log1p(df_tile.runtime_mean)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for s in stats_list:\n    fig = plt.figure(figsize=(12, 6))\n    gs = gridspec.GridSpec(1, 2)\n    for i, (optim, df_) in enumerate(zip(OPTIM, [df_layout, df_tile])):\n        jg, title = _plot_bivar(df_, [s, \"runtime_mean_log\"])\n        _ = SeabornFig2Grid(jg, fig, gs[i])\n\n        # Add title\n        ax_ghost = fig.add_subplot(gs[i])\n        ax_ghost.axis(\"off\")\n        ax_ghost.set_title(f\"{optim} - {title}\")\n    plt.tight_layout()","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Same as what we observe in the plots above, the correlation heatmap shows that graph stats, excluding `avg_clust_coeff`, generally have positive correlation with average runtime. Also, strong positive correlations exist between multiple graph stats pairs.","metadata":{}},{"cell_type":"code","source":"fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(16, 6))\nfor i, (optim, df_) in enumerate(zip(OPTIM, [df_layout, df_tile])):\n    df_corr = df_[stats_list + [\"runtime_mean_log\"]].corr()\n    mask = np.triu(np.ones_like(df_corr, dtype=np.bool_))\n    sns.heatmap(df_corr, mask=mask, cmap=\"crest\", annot=True, ax=axes[i])\n    axes[i].set_title(f\"{optim} - Correlation of Graph Stats with (Log) Average Runtime\")\n    plt.tight_layout()","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"<a id=\"ref\"></a>\n## Reference\n[**<span style=\"color:#FEF1FE; background-color:#535d70;border-radius: 5px; padding: 2px\">Go to Table of Content</span>**](#toc)\n\n[1] [TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs](https://arxiv.org/pdf/2308.13490.pdf)<br>\n[2] [Topological sorting](https://en.wikipedia.org/wiki/Topological_sorting)","metadata":{}},{"cell_type":"markdown","source":"Thanks for your attention, please stay tuned to see more analysis and explanations!","metadata":{}}],"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}}