workspace/run_all.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "08bb5a11",
   "metadata": {},
   "source": [
    "## Run SecureLoop Scheduling\n",
    "\n",
    "Run all three steps of scheduling and compare the effect of different scheduling algorithms (Fig. 11). For more detailed experiments for each step, check 1) `run_loopnest_scheduling.ipynb`, 2) `run_authblock_assignment.ipynb`, and 3) `run_simulated_annealing.ipynb`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a945c8ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import os\n",
    "import yaml\n",
    "import shutil\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ea00bf7",
   "metadata": {},
   "source": [
    "### 1. Define the accelerator architecture / parameters\n",
    "\n",
    "First, define an architecture design. The code below generates/detects a new architecture configuration based on the template design at `designs/{design_name}/template`. We provide `eyeriss_like` architecture we used for experiments in the paper in the source code. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bf6ce96",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import generate_arch_files, xml2mapping \n",
    "\n",
    "configuration_dict = {}\n",
    "\n",
    "# template design (with constraints and memory hierarchy representing \"dataflow\")\n",
    "configuration_dict['TEMPLATE_DESIGN'] = 'eyeriss_like'\n",
    "\n",
    "# number of bits used for I/O/W; we assume integer\n",
    "configuration_dict['WORDBITS'] = 16\n",
    "\n",
    "# DRAM bandwidth setting: words / cycle (not bits / cycle)\n",
    "configuration_dict['DRAM_READ_BANDWIDTH'] = 32\n",
    "configuration_dict['DRAM_WRITE_BANDWIDTH'] = 32\n",
    "\n",
    "# SRAM setting\n",
    "# - do we have a single shared glb or multiple glbs for each datatype? \n",
    "# - for each glb (if shared, just one), define depth/width/#banks and bandwidths\n",
    "configuration_dict['SRAM_SHARED'] = True\n",
    "configuration_dict['SRAM_DEPTH'] = [2 ** 13]\n",
    "configuration_dict['SRAM_WIDTH'] = [2 ** 7]\n",
    "configuration_dict['SRAM_BANKS'] = [32]                     # SRAM width and SRAM banks define the maximum possible bandwidth\n",
    "configuration_dict['SRAM_READ_BANDWIDTH'] = [32]\n",
    "configuration_dict['SRAM_WRITE_BANDWIDTH'] = [32]\n",
    "\n",
    "# PE array setting\n",
    "# - shape of PE array X x Y\n",
    "# - whether a PE has a shared scratchpad or separate scratchpads for each datatype\n",
    "configuration_dict['PE_X'] = 14\n",
    "configuration_dict['PE_Y'] = 12\n",
    "configuration_dict['PE_SPAD_SHARED'] = False\n",
    "configuration_dict['PE_SPAD_DEPTH'] = [192, 12, 16]         # Weight, IFmap, OFmap\n",
    "configuration_dict['PE_SPAD_WIDTH'] = [16, 16, 16]\n",
    "\n",
    "# Cryptographic engine setting\n",
    "# - type of cryptographic engine + dram (LPDDR4 + AES-GCM)\n",
    "# - cycle for AES-GCM \n",
    "# - whether the cryptographic engines are shared among all datatypes or assigned to each datatype\n",
    "configuration_dict['CRYPT_ENGINE_TYPE'] = 'effective_lpddr4_aesgcm'\n",
    "configuration_dict['CRYPT_ENGINE_CYCLE_PER_BLOCK'] = 11            # avg. cycle/128bit\n",
    "\n",
    "configuration_dict['CRYPT_ENGINE_SHARED'] = False\n",
    "configuration_dict['CRYPT_ENGINE_COUNT'] = [1, 1, 1]\n",
    "\n",
    "configuration_dict['EFFECTIVE_CONSERVATIVE'] = True\n",
    "\n",
    "# Create directory for this configuration if it doesn't exist already\n",
    "# iterate through design folders to check if any pre-exisiting folder\n",
    "design_dir = 'designs/{}'.format(configuration_dict['TEMPLATE_DESIGN'])\n",
    "arch_dir = None\n",
    "total_vers = 0\n",
    "for path in os.listdir(design_dir):\n",
    "    if path != 'template' and os.path.isdir(os.path.join(design_dir, path)):\n",
    "        try:\n",
    "            with open(os.path.join(design_dir, path, 'config.yaml'), 'r') as f:\n",
    "                config_file = yaml.safe_load(f)\n",
    "            total_vers += 1\n",
    "            if config_file == configuration_dict:\n",
    "                arch_dir = path\n",
    "                print(\"Pre-existing folder found. Setting the arch_dir to {}\".format(arch_dir))\n",
    "                break\n",
    "        except:\n",
    "            print(\"No config.yaml file in the directory {}\".format(str(os.path.join(design_dir, path))))\n",
    "            \n",
    "if arch_dir is None:\n",
    "    arch_dir = 'ver{}'.format(total_vers)\n",
    "    shutil.copytree(os.path.join(design_dir, 'template'), os.path.join(design_dir, arch_dir))\n",
    "    with open(os.path.join(design_dir, arch_dir, 'config.yaml'), 'w') as f:\n",
    "        _ = yaml.dump(configuration_dict, f)\n",
    "    \n",
    "    # create baseline and effective files\n",
    "    generate_arch_files(os.path.join(design_dir, arch_dir, 'arch'), configuration_dict)\n",
    "    \n",
    "    # create scheduling / evaluation folder\n",
    "    os.mkdir(os.path.join(design_dir, arch_dir, 'scheduling'))\n",
    "    os.mkdir(os.path.join(design_dir, arch_dir, 'evaluation'))\n",
    "    \n",
    "    # create folders for baseline scheduling / evaluation\n",
    "    os.mkdir(os.path.join(design_dir, arch_dir, 'baseline_scheduling'))\n",
    "    os.mkdir(os.path.join(design_dir, arch_dir, 'baseline_evaluation'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bffd3420",
   "metadata": {},
   "source": [
    "### 2. Define the DNN workload\n",
    "\n",
    "Define a workload as PyTorch's torch.nn.module. Then, convert the workload to the Timeloop's workload format and extract unique layers and interlayer dependency information."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50f55fa9",
   "metadata": {},
   "source": [
    "#### 2-1. Convert PyTorch workload to Timeloop workload (< 1 min)\n",
    "\n",
    "Define your workload. We used AlexNet, Resnet18, and MobilenetV2 for experiments. **Comment in/out the model you want to test, or define your own model here.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55800535",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import torchvision.models as model_zoo\n",
    "\n",
    "import pytorch2timeloop as pytorch2timeloop\n",
    "\n",
    "# Note: this version only supports nn.Conv2d (both normal convs and depthwise/pointwise convs) and nn.Linear\n",
    "\n",
    "# AlexNet\n",
    "# model_name = 'alexnet'\n",
    "# net = model_zoo.alexnet(pretrained=False)\n",
    "# layers_exclude_from_search = [6, 7, 8] # Define layer idx if you don't want to search them for simulated anneling (e.g., non-conv layers in AlexNet)\n",
    "\n",
    "# ResNet18\n",
    "# model_name = 'resnet18'\n",
    "# net = model_zoo.resnet18(pretrained=False)\n",
    "# layers_exclude_from_search = []\n",
    "\n",
    "# MobilenetV2\n",
    "model_name = 'mobilenet_v2'\n",
    "net = model_zoo.mobilenet_v2(pretrained=False)\n",
    "layers_exclude_from_search = []\n",
    "\n",
    "# Input / Batch info\n",
    "input_size = (3, 224, 224)\n",
    "batch_size = 1\n",
    "\n",
    "# print(net)\n",
    "\n",
    "# Convert to timeloop workloads; stored in workloads/{model_name}_batch{batch_size}\n",
    "top_dir = 'workloads'\n",
    "sub_dir = '{}_batch{}'.format(model_name, batch_size)\n",
    "exception_module_names = []\n",
    "\n",
    "overwrite = False\n",
    "if not os.path.exists(os.path.join(top_dir, sub_dir)) or overwrite:\n",
    "    pytorch2timeloop.convert_model(\n",
    "            net,\n",
    "            input_size,\n",
    "            batch_size,\n",
    "            sub_dir,\n",
    "            top_dir,\n",
    "            True,\n",
    "            exception_module_names\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb6fde1a",
   "metadata": {},
   "source": [
    "#### 2-2. Extract unique layers (per-layer Timeloop loopnest scheduling only for unqiue layers) (< 1 min)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca839524",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_dir = Path(os.getcwd())\n",
    "timeloop_dir = 'designs/{}/{}'.format(configuration_dict['TEMPLATE_DESIGN'], arch_dir)\n",
    "\n",
    "n_layers = 0\n",
    "layer_dict = {}\n",
    "layer_duplicate_info = {}\n",
    "unique_layers = []\n",
    "for module in net.modules():\n",
    "    if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):\n",
    "        n_layers += 1\n",
    "        if n_layers not in layer_dict.keys():\n",
    "            workload_path = os.path.join(base_dir, top_dir, sub_dir, '{}_layer{}.yaml'.format(sub_dir, n_layers))\n",
    "            with open(workload_path, 'r') as f:\n",
    "                workload_info = yaml.safe_load(f)\n",
    "            layer_dict[n_layers] = workload_info\n",
    "        \n",
    "        # identify the earliest duplicate layer\n",
    "        for key in range(1, n_layers):\n",
    "            if layer_dict[key] == layer_dict[n_layers]:\n",
    "                layer_duplicate_info[n_layers] = key\n",
    "                break\n",
    "        if n_layers not in layer_duplicate_info:\n",
    "            unique_layers.append(n_layers)\n",
    "            \n",
    "# print(layer_duplicate_info)\n",
    "# print(unique_layers)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7edcbb80",
   "metadata": {},
   "source": [
    "#### 2-3. Extract interlayer dependency (1~10 min if running for the first time; < 1 min if saved yaml file is available)\n",
    "\n",
    "We need to take care of interlayer dependencies. Since the layer number does not imply the actual layer order/dependency, and often multiple layers can be dependent on one layer and vice versa (e.g., residual connections), we create a back-propagation graph and determine the dependency from the graph. \n",
    "\n",
    "Depending on your accelerator architecture assumptions, some post-processing operations can be performed on-the-fly. Here, we assume that ReLU activation and batch normalization can be performed on-the-fly (thus, not affect the dependency), while pooling operations and adding multiple feature maps together (e.g., adding residual branches together) cannot be done on-the-fly (thus, break the depenency).\n",
    "\n",
    "If you have different assumptions about the post-processing or using different versions of PyTorch with different backend operations, `BackpropGraph.isDependentLayer` has to be modified accordingly.\n",
    "\n",
    "Define whether you want to ignore the interlayer dependency entirely (i.e., always use rehashing). Then, the dependency dictionary will ignore all interlayer dependency. \n",
    "\n",
    "Finally, constructing and analyzing the back-propagation grpah can take > 5 minutes for deep models like MobilenetV2. Since this dictionary has to be constructed only once per each model, we first search for existing dictionary. If not, then we generate a graph and create a dictionary and save it. \n",
    "\n",
    "*We provide the layer_info.yaml files in `workloads` - if you want to generate your own info, remove those yaml files*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8e9b675",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_layer_dependency_utils import BackpropGraph\n",
    "\n",
    "workload_path_1 = os.path.join(base_dir, top_dir, sub_dir, 'layer_info_interlayer.yaml')\n",
    "workload_path_2 = os.path.join(base_dir, top_dir, sub_dir, 'layer_info_ignore_interlayer.yaml') # we also need this for baseline\n",
    "\n",
    "try:\n",
    "    with open(workload_path_1, 'r') as f:\n",
    "        layer_info = yaml.safe_load(f)\n",
    "    with open(workload_path_2, 'r') as f:\n",
    "        layer_info_ignore_interlayer = yaml.safe_load(f)\n",
    "        \n",
    "except:\n",
    "\n",
    "    graph = BackpropGraph(net, [1, input_size[0], input_size[1], input_size[2]])\n",
    "    consecutive_dict, dependent_dict = graph.get_dependency_info()\n",
    "\n",
    "    # construct layer_info\n",
    "    layer_info = {}\n",
    "    layer_info_ignore_interlayer = {}\n",
    "    for layer_idx in range(1, n_layers + 1):\n",
    "        info = {}\n",
    "        if layer_idx in unique_layers:\n",
    "            info['layer_id_for_timeloop'] = layer_idx\n",
    "        else:\n",
    "            info['layer_id_for_timeloop'] = layer_duplicate_info[layer_idx]\n",
    "        info['prev_layer'] = []\n",
    "        info['next_layer'] = []\n",
    "        info['dependent_prev_layer'] = []\n",
    "        info['dependent_next_layer'] = []\n",
    "        layer_info[layer_idx] = info\n",
    "        layer_info_ignore_interlayer[layer_idx] = info\n",
    "\n",
    "    for layer_idx in range(1, n_layers + 1):\n",
    "        consecutive = consecutive_dict[layer_idx]\n",
    "        dependent = dependent_dict[layer_idx]\n",
    "        layer_info[layer_idx]['next_layer'].extend(consecutive)\n",
    "        layer_info_ignore_interlayer[layer_idx]['next_layer'].extend(consecutive)\n",
    "        for i in consecutive:\n",
    "            layer_info[i]['prev_layer'].append(layer_idx)\n",
    "            layer_info_ignore_interlayer[i]['prev_layer'].append(layer_idx)\n",
    "        if len(dependent) > 0 and not ignore_interlayer:\n",
    "            layer_info[layer_idx]['dependent_next_layer'].extend(dependent)   \n",
    "            for i in dependent:\n",
    "                layer_info[i]['dependent_prev_layer'].append(layer_idx)\n",
    "\n",
    "    # for layer_idx in range(1, n_layers + 1):\n",
    "    #     print(layer_idx, layer_info[layer_idx])\n",
    "\n",
    "    # store therresults - this can take long for deep models like MobileNet..\n",
    "    # f\"{base_dir/top_dir/sub_dir/sub_dir}_layer{layer_id}.yaml \"\n",
    "    with open(workload_path_1, 'w') as f:\n",
    "        _ = yaml.dump(layer_info, f)\n",
    "    with open(workload_path_2, 'w') as f:\n",
    "        _ = yaml.dump(layer_info_ignore_interlayer, f)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "225e1513",
   "metadata": {},
   "source": [
    "### 3. Define the top-k parameter and run timeloop-topk using the effective model (Step 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "043e11ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "topk = 6\n",
    "mapper_file_path = os.path.join(base_dir, timeloop_dir, 'mapper/mapper.yaml')\n",
    "with open(mapper_file_path, 'r') as f:\n",
    "    mapper_config = yaml.safe_load(f)\n",
    "mapper_config['mapper']['topk'] = topk\n",
    "with open(mapper_file_path, 'w') as f:\n",
    "    _ = yaml.dump(mapper_config, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dd1b2ac",
   "metadata": {},
   "source": [
    "#### 3-1. Run timeloop-topk for all unique layers (~30 min)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77665669",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def get_cmd(workload_info, layer_id, base_dir, timeloop_dir, sub_dir, top_dir):\n",
    "    cwd = f\"{base_dir/timeloop_dir/'scheduling'/sub_dir/f'layer{layer_id}'}\"\n",
    "    if 'M' in workload_info['problem']['instance']:\n",
    "        constraint_pth = base_dir/timeloop_dir/'constraints/*.yaml'\n",
    "    else:\n",
    "        # depthwise\n",
    "        constraint_pth = base_dir/timeloop_dir/'constraints_dw/*.yaml'\n",
    "\n",
    "    timeloopcmd = f\"timeloop-mapper-topk \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'arch/effective.yaml'} \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'arch/components/*.yaml'} \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'mapper/mapper.yaml'} \" \\\n",
    "                  f\"{constraint_pth} \" \\\n",
    "                  f\"{base_dir/top_dir/sub_dir/sub_dir}_layer{layer_id}.yaml \"\n",
    "    return [cwd, timeloopcmd]\n",
    "\n",
    "cwd_list = []\n",
    "cmd_list = []\n",
    "\n",
    "for layer_id in unique_layers:\n",
    "    workload_path = os.path.join(base_dir, top_dir, sub_dir, '{}_layer{}.yaml'.format(sub_dir, layer_id))\n",
    "    with open(workload_path, 'r') as f:\n",
    "        workload_info = yaml.safe_load(f)\n",
    "    [cwd, cmd] = get_cmd(workload_info, layer_id, base_dir, timeloop_dir, sub_dir, top_dir)\n",
    "    cwd_list.append(cwd)\n",
    "    cmd_list.append(cmd)\n",
    "    \n",
    "if not os.path.exists(os.path.join(base_dir, timeloop_dir, 'scheduling', sub_dir)):\n",
    "    os.mkdir(os.path.join(base_dir, timeloop_dir, 'scheduling', sub_dir))\n",
    "    \n",
    "start_time = time.time()\n",
    "for cwd, cmd in zip(cwd_list, cmd_list):\n",
    "    print(\"Executing cmd: {}\".format(cmd))\n",
    "    try:\n",
    "        os.chdir(cwd)\n",
    "    except:\n",
    "        os.mkdir(cwd)\n",
    "        os.chdir(cwd)\n",
    "    os.system(cmd)\n",
    "os.chdir(base_dir)\n",
    "\n",
    "# Time this cell\n",
    "print(\"Execution time: {}s\".format(time.time() - start_time))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d96735c7",
   "metadata": {},
   "source": [
    "#### 3-2. Convert the found schedule to yaml files (~ 1 min)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed55610",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_to_mapping(base_dir, timeloop_dir, top_dir, sub_dir, layer_idx, topk_idx):\n",
    "    xml_file = os.path.join(base_dir, timeloop_dir, 'scheduling', sub_dir, \"layer{}\".format(layer_idx), \\\n",
    "                            \"timeloop-mapper-topk{}.map+stats.xml\".format(topk_idx))\n",
    "    workload_file = os.path.join(base_dir, top_dir, sub_dir, \"{}_layer{}.yaml\".format(sub_dir, layer_idx))\n",
    "    # print(workload_file)\n",
    "    with open(workload_file, 'r') as f:\n",
    "        workload_info = yaml.safe_load(f)\n",
    "    if 'M' in workload_info['problem']['instance']:\n",
    "        dw = False\n",
    "    else:\n",
    "        dw = True\n",
    "    arch_constraint_file = os.path.join(base_dir, timeloop_dir, 'constraints_dw' if dw else 'constraints' , \\\n",
    "                                        'eyeriss_like_arch_constraints.yaml' if (configuration_dict['TEMPLATE_DESIGN'] == 'eyeriss_like' \\\n",
    "                                                                                 or configuration_dict['TEMPLATE_DESIGN'] == 'eyeriss_like_hbm2') \\\n",
    "                                        else 'simple_output_stationary_arch_constraints.yaml' if configuration_dict['TEMPLATE_DESIGN'] == 'output_stationary' \\\n",
    "                                        else 'simple_weight_stationary_arch_constraints.yaml')\n",
    "    # print(layer_idx, dw)\n",
    "    mapping = xml2mapping(xml_file, workload_file, arch_constraint_file, dw)\n",
    "    with open(os.path.join(base_dir, timeloop_dir, 'scheduling',sub_dir, \"layer{}\".format(layer_idx), \\\n",
    "                           \"mapping{}.yaml\".format(topk_idx)), 'w') as f:\n",
    "        _ = yaml.dump({'mapping': mapping}, f)\n",
    "        \n",
    "for layer_idx in unique_layers:\n",
    "    for k in range(1, topk + 1):\n",
    "        convert_to_mapping(base_dir, timeloop_dir, top_dir, sub_dir, layer_idx, k)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67395b62",
   "metadata": {},
   "source": [
    "#### 3-3. Evaluate the top-1 loopnest schedule (< 5 min)\n",
    "\n",
    "Note that we use the effective off-chip bandwidth model for **scheduling**, but the actual energy and latency of the accelerator (excluding the cryptographic engine) has to evaluated with the baseline model. The actual cost of cryptographic operations is more complex (i.e., AuthBlock assignment) and is added in the next step. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02840733",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_cmd_model(workload_info, layer_id, base_dir, timeloop_dir, sub_dir, top_dir):\n",
    "    cwd = f\"{base_dir/timeloop_dir/'evaluation'/sub_dir/f'layer{layer_id}'}\"\n",
    "    if 'M' in workload_info['problem']['instance']:\n",
    "        constraint_pth = base_dir/timeloop_dir/'constraints/*.yaml'\n",
    "    else:\n",
    "        # depthwise\n",
    "        constraint_pth = base_dir/timeloop_dir/'constraints_dw/*.yaml'\n",
    "\n",
    "    timeloopcmd = f\"timeloop-model \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'arch/baseline.yaml'} \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'arch/components/*.yaml'} \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'scheduling'/sub_dir/f'layer{layer_id}/mapping1.yaml'} \" \\\n",
    "                  f\"{base_dir/top_dir/sub_dir/sub_dir}_layer{layer_id}.yaml \"\n",
    "    return [cwd, timeloopcmd]\n",
    "\n",
    "cwd_list = []\n",
    "cmd_list = []\n",
    "for layer_id in unique_layers:\n",
    "    workload_path = os.path.join(base_dir, top_dir, sub_dir, '{}_layer{}.yaml'.format(sub_dir, layer_id))\n",
    "    with open(workload_path, 'r') as f:\n",
    "        workload_info = yaml.safe_load(f)\n",
    "    [cwd, cmd] = get_cmd_model(workload_info, layer_id, base_dir, timeloop_dir, sub_dir, top_dir)\n",
    "    cwd_list.append(cwd)\n",
    "    cmd_list.append(cmd)\n",
    "    \n",
    "if not os.path.exists(os.path.join(base_dir, timeloop_dir, 'evaluation', sub_dir)):\n",
    "    os.mkdir(os.path.join(base_dir, timeloop_dir, 'evaluation', sub_dir))\n",
    "for cwd, cmd in zip(cwd_list, cmd_list):\n",
    "    print(\"Executing cmd: {}\".format(cmd))\n",
    "    try:\n",
    "        os.chdir(cwd)\n",
    "    except:\n",
    "        os.mkdir(cwd)\n",
    "        os.chdir(cwd)\n",
    "    os.system(cmd)\n",
    "os.chdir(base_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd175a12",
   "metadata": {},
   "source": [
    "### 4. AuthBlock Assignment (tile-as-an-AuthBlock vs. our algorithm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b98aeb0",
   "metadata": {},
   "source": [
    "#### 4-1. tile-as-an-AuthBlock using the per-layer top-1 loopnest schedule (*Crypt-Tile-Single*) (< 5 min)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f33d335e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from authblock_assignment import AuthBlockAssignment\n",
    "\n",
    "cts_cost_dict, cts_rehash_cost_dict, cts_block_info_dict = \\\n",
    "AuthBlockAssignment(n_layers, layer_info_ignore_interlayer, \\\n",
    "                    base_dir, timeloop_dir, top_dir, sub_dir, \\\n",
    "                    configuration_dict, mode=\"tile\", \\\n",
    "                    joint=False, generate_summary=True, return_cost_dict=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "097a12dd",
   "metadata": {},
   "source": [
    "#### 4-2. Our optimal AuthBlock assignment algorithm using the per-layer top-1 loopnest schedule (*Crypt-Opt-Single*) (< 5 min)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e5c0782",
   "metadata": {},
   "outputs": [],
   "source": [
    "from authblock_assignment import AuthBlockAssignment\n",
    "\n",
    "cos_cost_dict, cos_rehash_cost_dict, cos_block_info_dict = \\\n",
    "AuthBlockAssignment(n_layers, layer_info, \\\n",
    "                    base_dir, timeloop_dir, top_dir, sub_dir, \\\n",
    "                    configuration_dict, mode=\"search\", \\\n",
    "                    joint=False, generate_summary=True, return_cost_dict=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3094fc9a",
   "metadata": {},
   "source": [
    "### 5. Simulated annealing for interlayer dependency \n",
    "\n",
    "Run simulated annealing to identify the optimal loopnest schedule when multiple layers are jointly explored. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fcc693d",
   "metadata": {},
   "source": [
    "#### 5-1. Prepare folders before running simulated annealing (< 5 min)\n",
    "\n",
    "We have to evaluate all top-k loopnest schedules for reference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4bf45e1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from authblock_assignment import AuthBlockAssignment\n",
    "\n",
    "if not os.path.exists(os.path.join(base_dir, timeloop_dir, 'joint_topk')):\n",
    "    os.mkdir(os.path.join(base_dir, timeloop_dir, 'joint_topk'))\n",
    "if not os.path.exists(os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir)):\n",
    "    os.mkdir(os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir))\n",
    "    \n",
    "base_cost_dict, base_rehash_cost_dict, base_block_info_dict = AuthBlockAssignment(n_layers, layer_info, \\\n",
    "                                                                                  base_dir, timeloop_dir, top_dir, sub_dir, \\\n",
    "                                                                                  configuration_dict, \\\n",
    "                                                                                  mode=\"search\", \\\n",
    "                                                                                  joint=False, return_cost_dict=True)\n",
    "\n",
    "baseline_energy = 0\n",
    "baseline_latency = 0\n",
    "baseline_add_mem_traffic = 0\n",
    "\n",
    "for key in base_cost_dict:\n",
    "    baseline_energy += base_cost_dict[key]['total_energy'] / 10**6\n",
    "    baseline_latency += base_cost_dict[key]['total_latency']\n",
    "    baseline_add_mem_traffic += base_cost_dict[key]['add_memory_traffic']\n",
    "for key in base_rehash_cost_dict:\n",
    "    baseline_energy += base_rehash_cost_dict[key]['total_energy'] / 10**6\n",
    "    baseline_latency += base_rehash_cost_dict[key]['total_latency']\n",
    "    baseline_add_mem_traffic += base_rehash_cost_dict[key]['add_memory_traffic']   \n",
    "    \n",
    "for layer_idx in range(1, n_layers + 1):\n",
    "    work_dir = os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir, 'layer{}'.format(layer_idx))\n",
    "    if not os.path.exists(work_dir):\n",
    "        os.mkdir(work_dir)\n",
    "        \n",
    "    # \"\"\"\n",
    "    for k in range(1, topk + 1):\n",
    "        if not os.path.exists(os.path.join(work_dir, 'eval{}'.format(k))):\n",
    "            os.mkdir(os.path.join(work_dir, 'eval{}'.format(k)))\n",
    "        layer_id_for_timeloop = layer_info[layer_idx]['layer_id_for_timeloop']\n",
    "        cwd = f\"{base_dir/timeloop_dir/'joint_topk'/sub_dir/f'layer{layer_idx}'/f'eval{k}'}\"\n",
    "        if 'M' in workload_info['problem']['instance']:\n",
    "            constraint_pth = base_dir/timeloop_dir/'constraints/*.yaml'\n",
    "        else:\n",
    "            # depthwise\n",
    "            constraint_pth = base_dir/timeloop_dir/'constraints_dw/*.yaml'\n",
    "\n",
    "        timeloopcmd = f\"timeloop-model \" \\\n",
    "              f\"{base_dir/timeloop_dir/'arch/baseline.yaml'} \" \\\n",
    "              f\"{base_dir/timeloop_dir/'arch/components/*.yaml'} \" \\\n",
    "              f\"{base_dir/timeloop_dir/'scheduling'/sub_dir/f'layer{layer_id_for_timeloop}'/f'mapping{k}.yaml'} \" \\\n",
    "              f\"{base_dir/top_dir/sub_dir/sub_dir}_layer{layer_idx}.yaml \"\n",
    "        \n",
    "        try:\n",
    "            os.chdir(cwd)\n",
    "        except:\n",
    "            os.mkdir(cwd)\n",
    "            os.chdir(cwd)\n",
    "        os.system(timeloopcmd)\n",
    "        os.chdir(base_dir)\n",
    "    # \"\"\"\n",
    "\n",
    "    # copy mapping1's result into here\n",
    "    shutil.copy(os.path.join(work_dir, 'eval1', 'timeloop-model.map+stats.xml'), work_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "324d95fe",
   "metadata": {},
   "source": [
    "#### 5-2. Run simulated annealing (~ 10 min for AlexNet/MobilenetV2; ~90 min ResNet18)\n",
    "\n",
    "First, define the hyperparameters for simulated annealing. Then, run the algorithm.\n",
    "\n",
    "For every iteration of simulated annealing, AuthBlockAssignment has to be executed to evaluate the cost. To make the search process faster, we support partial update of the AuthBlockAssignment (only calculate for changed layer and its dependent layers). However, there is currently some bugs with supporting the partial update for ResNet18 and we run the full AuthBlockAssignment for ResNet18 instead. As such, this cell takes much longer to execute for ResNet18. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba1ee0e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_temp = 100\n",
    "final_temp = 0.1\n",
    "n_iters = 1000\n",
    "\n",
    "cooling_scheduler = 'linear'\n",
    "\n",
    "# TODO: this option should not be used for ResNet18 - bug with dependent layer partial update due to residuals\n",
    "use_partial_update = True\n",
    "if model_name == 'resnet18':\n",
    "    use_partial_update = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d046093",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import time\n",
    "import csv\n",
    "import copy\n",
    "import math\n",
    "\n",
    "from authblock_assignment import PartialUpdateAuthBlockAssignment\n",
    "\n",
    "csv_header = ['Iter', 'Temp', \\\n",
    "              'Cost (J x cycles)', 'Total Latency (cycles)', 'Total Energy (uJ)', 'Additional Off-chip Traffic (bits)']\n",
    "logs = []\n",
    "\n",
    "solution_cost_dict = copy.deepcopy(base_cost_dict)\n",
    "solution_rehash_cost_dict = copy.deepcopy(base_rehash_cost_dict)\n",
    "solution_block_info_dict = copy.deepcopy(base_block_info_dict)\n",
    "\n",
    "current_cost_dict = copy.deepcopy(base_cost_dict)\n",
    "current_rehash_cost_dict = copy.deepcopy(base_rehash_cost_dict)\n",
    "current_block_info_dict = copy.deepcopy(base_block_info_dict)\n",
    "\n",
    "solution_state = [1] * n_layers\n",
    "current_state = [1] * n_layers\n",
    "best_state = [1] * n_layers\n",
    "\n",
    "i = 0\n",
    "cost_best = baseline_energy * baseline_latency\n",
    "\n",
    "layers_for_search = []\n",
    "for idx in range(1, n_layers + 1):\n",
    "    if len(layer_info[idx]['dependent_next_layer']) > 0 or len(layer_info[idx]['dependent_prev_layer']) > 0:\n",
    "        if idx not in layers_exclude_from_search:\n",
    "            layers_for_search.append(idx)\n",
    "            \n",
    "start_time = time.time()\n",
    "while i < n_iters + 1:\n",
    "    # temperature\n",
    "    if cooling_scheduler == 'linear':\n",
    "        current_temp = final_temp + (initial_temp - final_temp) / float(n_iters) * float(n_iters - i)\n",
    "    elif cooling_scheduler == 'cosine':\n",
    "        current_temp = final_temp + 0.5 * (initial_temp - final_temp) * (1 + math.cos(float(i) * math.pi / float(n_iters)))\n",
    "    elif cooling_scheduler == 'quadratic':\n",
    "        current_temp = final_temp + (initial_temp - final_temp) * (float(n_iters - i) / float(n_iters))**2\n",
    "    \n",
    "    layer2change = random.choice(layers_for_search)\n",
    "    neighbor_loopnest = random.choice(list(range(1, topk + 1)))\n",
    "    \n",
    "    current_state[layer2change - 1] = neighbor_loopnest\n",
    "    stats_file = os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir, \"layer{}\".format(layer2change), \\\n",
    "                              \"eval{}\".format(neighbor_loopnest), \"timeloop-model.stats.txt\")\n",
    "    with open(stats_file, 'r') as f:\n",
    "        lines = f.read().split('\\n')[-200:]\n",
    "        for line in lines:\n",
    "            if line.startswith('Energy'):\n",
    "                energy = eval(line.split(': ')[1].split(' ')[0]) * float(10**6) # micro to pico\n",
    "                # print(energy)\n",
    "            elif line.startswith('Cycles'):\n",
    "                cycle = eval(line.split(': ')[1])\n",
    "    current_cost_dict[layer2change]['timeloop_energy'] = energy\n",
    "    current_cost_dict[layer2change]['timeloop_cycle'] = cycle\n",
    "    \n",
    "    xml_file = os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir, \"layer{}\".format(layer2change), \\\n",
    "                            \"eval{}\".format(neighbor_loopnest), \"timeloop-model.map+stats.xml\")\n",
    "    shutil.copy(xml_file, os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir, 'layer{}'.format(layer2change)))\n",
    "    \n",
    "    if use_partial_update:\n",
    "        subset_layers = [layer2change]\n",
    "        subset_layers.extend(layer_info[layer2change]['prev_layer'])\n",
    "        subset_layers.extend(layer_info[layer2change]['next_layer'])\n",
    "        \n",
    "        current_cost_dict, current_rehash_cost_dict, current_block_info_dict = \\\n",
    "        PartialUpdateAuthBlockAssignment(n_layers, layer_info, \\\n",
    "                                         base_dir, timeloop_dir, top_dir, sub_dir, \\\n",
    "                                         configuration_dict, mode=\"search\", \\\n",
    "                                         prev_block_info_dict=current_block_info_dict, subset_layers=subset_layers, \\\n",
    "                                         prev_cost_dict=current_cost_dict, prev_rehash_cost_dict=current_rehash_cost_dict)\n",
    "        \n",
    "    else:\n",
    "        current_cost_dict, current_rehash_cost_dict, current_block_info_dict = \\\n",
    "        PartialUpdateAuthBlockAssignment(n_layers, layer_info, \\\n",
    "                                         base_dir, timeloop_dir, top_dir, sub_dir, \\\n",
    "                                         configuration_dict, \\\n",
    "                                         mode=\"search\", \\\n",
    "                                         prev_block_info_dict=None, subset_layers=[], \\\n",
    "                                         prev_cost_dict=current_cost_dict, prev_rehash_cost_dict=None)\n",
    "        \n",
    "    solution_energy, solution_latency, solution_add_mem_traffic = 0, 0, 0\n",
    "    for key in solution_cost_dict:\n",
    "        solution_energy += solution_cost_dict[key]['total_energy'] / 10**6\n",
    "        solution_latency += solution_cost_dict[key]['total_latency']\n",
    "        solution_add_mem_traffic += solution_cost_dict[key]['add_memory_traffic']\n",
    "    for key in solution_rehash_cost_dict:\n",
    "        solution_energy += solution_rehash_cost_dict[key]['total_energy'] / 10**6\n",
    "        solution_latency += solution_rehash_cost_dict[key]['total_latency']\n",
    "        solution_add_mem_traffic += solution_rehash_cost_dict[key]['add_memory_traffic']\n",
    "    \n",
    "    current_energy, current_latency, current_add_mem_traffic = 0, 0, 0\n",
    "    for key in current_cost_dict:\n",
    "        current_energy += current_cost_dict[key]['total_energy'] / 10**6\n",
    "        current_latency += current_cost_dict[key]['total_latency']\n",
    "        current_add_mem_traffic += current_cost_dict[key]['add_memory_traffic']\n",
    "    for key in current_rehash_cost_dict:\n",
    "        current_energy += current_rehash_cost_dict[key]['total_energy'] / 10**6\n",
    "        current_latency += current_rehash_cost_dict[key]['total_latency']\n",
    "        current_add_mem_traffic += current_rehash_cost_dict[key]['add_memory_traffic']\n",
    "    \n",
    "    cost_solution = solution_energy * solution_latency\n",
    "    cost_current = current_energy * current_latency\n",
    "    cost_diff = (cost_solution - cost_current) / (10 ** 6 * n_layers)\n",
    "    \n",
    "    if cost_current < cost_best:\n",
    "        best_state = copy.deepcopy(current_state)\n",
    "        cost_best = cost_current\n",
    "        print(\"Found best so far: \", best_state, \" .. updating cost_best: {}\".format(cost_best))\n",
    "        \n",
    "    if cost_diff > 0 or (random.uniform(0, 1) < math.exp(cost_diff / current_temp)):\n",
    "        solution_state = copy.deepcopy(current_state)\n",
    "        solution_cost_dict = copy.deepcopy(current_cost_dict)\n",
    "        solution_rehash_cost_dict = copy.deepcopy(current_rehash_cost_dict)\n",
    "        solution_block_info_dict = copy.deepcopy(current_block_info_dict)\n",
    "    else:\n",
    "        # roll-back to the solution state\n",
    "        xml_file = os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir, \"layer{}\".format(layer2change), \\\n",
    "                                  \"eval{}\".format(solution_state[layer2change - 1]), \"timeloop-model.map+stats.xml\")\n",
    "        shutil.copy(xml_file, os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir, 'layer{}'.format(layer2change)))\n",
    "        current_state = copy.deepcopy(solution_state)\n",
    "        current_cost_dict = copy.deepcopy(solution_cost_dict)\n",
    "        current_rehash_cost_dict = copy.deepcopy(solution_rehash_cost_dict)\n",
    "        current_block_info_dict = copy.deepcopy(solution_block_info_dict)\n",
    "    \n",
    "    solution_energy, solution_latency, solution_add_mem_traffic = 0, 0, 0\n",
    "    for key in solution_cost_dict:\n",
    "        solution_energy += solution_cost_dict[key]['total_energy'] / 10**6\n",
    "        solution_latency += solution_cost_dict[key]['total_latency']\n",
    "        solution_add_mem_traffic += solution_cost_dict[key]['add_memory_traffic']\n",
    "    for key in solution_rehash_cost_dict:\n",
    "        solution_energy += solution_rehash_cost_dict[key]['total_energy'] / 10**6\n",
    "        solution_latency += solution_rehash_cost_dict[key]['total_latency']\n",
    "        solution_add_mem_traffic += solution_rehash_cost_dict[key]['add_memory_traffic']\n",
    "        \n",
    "    # print(\"Solution state: \", solution_state)\n",
    "    print(\"Current iteration: {} (temperature: {:.2f}) -- Latency: {} ({:.2f}% faster), Energy: {} uW ({:.2f}% lower), Add Mem Traffic: {} bits ({:.2f}% smaller)\"\\\n",
    "          .format(i+1, current_temp, solution_latency, (baseline_latency - solution_latency) / float(baseline_latency) * 100. , \\\n",
    "                  solution_energy, (baseline_energy - solution_energy) / baseline_energy * 100., \\\n",
    "                  solution_add_mem_traffic, (baseline_add_mem_traffic - solution_add_mem_traffic) / float(baseline_add_mem_traffic) * 100.))\n",
    "\n",
    "    curr_log = [(i + 1), current_temp, cost_solution, solution_latency, solution_energy, solution_add_mem_traffic]\n",
    "    logs.append(curr_log)\n",
    "    i += 1\n",
    "    \n",
    "    if current_temp < final_temp:\n",
    "        break\n",
    "        \n",
    "print(\"Execution time: {}s\".format(time.time() - start_time))\n",
    "\n",
    "# dump to csv file\n",
    "with open(os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir, 'SA_{}_top{}_summary.csv'.format(cooling_scheduler, topk)), 'w', newline='') as f:\n",
    "    writer = csv.writer(f)\n",
    "    writer.writerow(csv_header)\n",
    "    writer.writerows(logs)\n",
    "    \n",
    "# dump best state & solution state to yaml file\n",
    "state = {'best': best_state, 'final': solution_state}\n",
    "with open(os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir, 'SA_{}_state.yaml'.format(cooling_scheduler)), 'w') as f:\n",
    "    _ = yaml.dump(state, f)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af7fe654",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir, 'SA_{}_state.yaml'.format('linear')), 'r') as f:\n",
    "    states = yaml.safe_load(f)\n",
    "    best_state = states['best']\n",
    "\n",
    "# move the best solution result\n",
    "for layer_idx in range(1, n_layers + 1):\n",
    "    loopnest_id = best_state[layer_idx - 1]\n",
    "    src = os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir, 'layer{}'.format(layer_idx), \\\n",
    "                       'eval{}'.format(loopnest_id))\n",
    "    src_files = os.listdir(src)\n",
    "    for file in src_files:\n",
    "        file_name = os.path.join(src, file)\n",
    "        if os.path.isfile(file_name):\n",
    "            shutil.copy(file_name, os.path.join(os.path.join(base_dir, timeloop_dir, 'joint_topk', sub_dir, 'layer{}'.format(layer_idx))))\n",
    "            "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "531170d7",
   "metadata": {},
   "source": [
    "#### 5-3. Evaluate with AuthBlock Assignment (*Crypt-Opt-Cross*) (< 5 min)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79f354b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from authblock_assignment import AuthBlockAssignment\n",
    "\n",
    "coc_cost_dict, coc_rehash_cost_dict, coc_block_info_dict = \\\n",
    "AuthBlockAssignment(n_layers, layer_info, \\\n",
    "                    base_dir, timeloop_dir, top_dir, sub_dir, \\\n",
    "                    configuration_dict, mode=\"search\", \\\n",
    "                    joint=True, generate_summary=True, return_cost_dict=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b898fc9",
   "metadata": {},
   "source": [
    "### 6. Putting all together: plot the graph (Fig. 11)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc747e72",
   "metadata": {},
   "source": [
    "#### 6-1. Get the unsecured baseline architecture's latency and energy (~ 30 min)\n",
    "\n",
    "To measure the slowdown over the unsecure design, we have to run timeloop-topk for the baseline architecture. Repeat the process in Section 3, but with the baseline architecture definition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8e8eb3b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_cmd(workload_info, layer_id, base_dir, timeloop_dir, sub_dir, top_dir):\n",
    "    cwd = f\"{base_dir/timeloop_dir/'baseline_scheduling'/sub_dir/f'layer{layer_id}'}\"\n",
    "    if 'M' in workload_info['problem']['instance']:\n",
    "        constraint_pth = base_dir/timeloop_dir/'constraints/*.yaml'\n",
    "    else:\n",
    "        # depthwise\n",
    "        constraint_pth = base_dir/timeloop_dir/'constraints_dw/*.yaml'\n",
    "\n",
    "    timeloopcmd = f\"timeloop-mapper-topk \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'arch/baseline.yaml'} \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'arch/components/*.yaml'} \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'mapper/mapper.yaml'} \" \\\n",
    "                  f\"{constraint_pth} \" \\\n",
    "                  f\"{base_dir/top_dir/sub_dir/sub_dir}_layer{layer_id}.yaml \"\n",
    "    return [cwd, timeloopcmd]\n",
    "\n",
    "cwd_list = []\n",
    "cmd_list = []\n",
    "\n",
    "for layer_id in unique_layers:\n",
    "    workload_path = os.path.join(base_dir, top_dir, sub_dir, '{}_layer{}.yaml'.format(sub_dir, layer_id))\n",
    "    with open(workload_path, 'r') as f:\n",
    "        workload_info = yaml.safe_load(f)\n",
    "    [cwd, cmd] = get_cmd(workload_info, layer_id, base_dir, timeloop_dir, sub_dir, top_dir)\n",
    "    cwd_list.append(cwd)\n",
    "    cmd_list.append(cmd)\n",
    "    \n",
    "if not os.path.exists(os.path.join(base_dir, timeloop_dir, 'baseline_scheduling', sub_dir)):\n",
    "    os.mkdir(os.path.join(base_dir, timeloop_dir, 'baseline_scheduling', sub_dir))\n",
    "for cwd, cmd in zip(cwd_list, cmd_list):\n",
    "    print(\"Executing cmd: {}\".format(cmd))\n",
    "    try:\n",
    "        os.chdir(cwd)\n",
    "    except:\n",
    "        os.mkdir(cwd)\n",
    "        os.chdir(cwd)\n",
    "    os.system(cmd)\n",
    "os.chdir(base_dir)\n",
    "\n",
    "def convert_to_mapping(base_dir, timeloop_dir, top_dir, sub_dir, layer_idx, topk_idx):\n",
    "    xml_file = os.path.join(base_dir, timeloop_dir, 'baseline_scheduling', sub_dir, \"layer{}\".format(layer_idx), \\\n",
    "                            \"timeloop-mapper-topk{}.map+stats.xml\".format(topk_idx))\n",
    "    workload_file = os.path.join(base_dir, top_dir, sub_dir, \"{}_layer{}.yaml\".format(sub_dir, layer_idx))\n",
    "    with open(workload_file, 'r') as f:\n",
    "        workload_info = yaml.safe_load(f)\n",
    "    if 'M' in workload_info['problem']['instance']:\n",
    "        dw = False\n",
    "    else:\n",
    "        dw = True\n",
    "    arch_constraint_file = os.path.join(base_dir, timeloop_dir, 'constraints_dw' if dw else 'constraints' , \\\n",
    "                                        'eyeriss_like_arch_constraints.yaml' if (configuration_dict['TEMPLATE_DESIGN'] == 'eyeriss_like' or \\\n",
    "                                                                                 configuration_dict['TEMPLATE_DESIGN'] == 'eyeriss_like_hbm2') \\\n",
    "                                        else 'simple_output_stationary_arch_constraints.yaml' if configuration_dict['TEMPLATE_DESIGN'] == 'output_stationary' \\\n",
    "                                        else 'simple_weight_stationary_arch_constraints.yaml')\n",
    "    mapping = xml2mapping(xml_file, workload_file, arch_constraint_file, dw)\n",
    "    with open(os.path.join(base_dir, timeloop_dir, 'baseline_scheduling',sub_dir, \"layer{}\".format(layer_idx), \\\n",
    "                           \"mapping{}.yaml\".format(topk_idx)), 'w') as f:\n",
    "        _ = yaml.dump({'mapping': mapping}, f)\n",
    "        \n",
    "for layer_idx in unique_layers:\n",
    "    for k in range(1, topk + 1):\n",
    "        convert_to_mapping(base_dir, timeloop_dir, top_dir, sub_dir, layer_idx, k)\n",
    "        \n",
    "def get_cmd_model(workload_info, layer_id, base_dir, timeloop_dir, sub_dir, top_dir):\n",
    "    cwd = f\"{base_dir/timeloop_dir/'baseline_evaluation'/sub_dir/f'layer{layer_id}'}\"\n",
    "    if 'M' in workload_info['problem']['instance']:\n",
    "        constraint_pth = base_dir/timeloop_dir/'constraints/*.yaml'\n",
    "    else:\n",
    "        # depthwise\n",
    "        constraint_pth = base_dir/timeloop_dir/'constraints_dw/*.yaml'\n",
    "\n",
    "    timeloopcmd = f\"timeloop-model \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'arch/baseline.yaml'} \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'arch/components/*.yaml'} \" \\\n",
    "                  f\"{base_dir/timeloop_dir/'baseline_scheduling'/sub_dir/f'layer{layer_id}/mapping1.yaml'} \" \\\n",
    "                  f\"{base_dir/top_dir/sub_dir/sub_dir}_layer{layer_id}.yaml \"\n",
    "    return [cwd, timeloopcmd]\n",
    "\n",
    "cwd_list = []\n",
    "cmd_list = []\n",
    "for layer_id in unique_layers:\n",
    "    workload_path = os.path.join(base_dir, top_dir, sub_dir, '{}_layer{}.yaml'.format(sub_dir, layer_id))\n",
    "    with open(workload_path, 'r') as f:\n",
    "        workload_info = yaml.safe_load(f)\n",
    "    [cwd, cmd] = get_cmd_model(workload_info, layer_id, base_dir, timeloop_dir, sub_dir, top_dir)\n",
    "    cwd_list.append(cwd)\n",
    "    cmd_list.append(cmd)\n",
    "    \n",
    "if not os.path.exists(os.path.join(base_dir, timeloop_dir, 'baseline_evaluation', sub_dir)):\n",
    "    os.mkdir(os.path.join(base_dir, timeloop_dir, 'baseline_evaluation', sub_dir))\n",
    "for cwd, cmd in zip(cwd_list, cmd_list):\n",
    "    print(\"Executing cmd: {}\".format(cmd))\n",
    "    try:\n",
    "        os.chdir(cwd)\n",
    "    except:\n",
    "        os.mkdir(cwd)\n",
    "        os.chdir(cwd)\n",
    "    os.system(cmd)\n",
    "os.chdir(base_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f980a16",
   "metadata": {},
   "outputs": [],
   "source": [
    "baseline_latency = 0\n",
    "baseline_energy = 0\n",
    "\n",
    "for layer_id in range(1, n_layers + 1):\n",
    "    if layer_id in layers_exclude_from_search:\n",
    "        continue\n",
    "    layer_id_for_timeloop = layer_info[layer_id]['layer_id_for_timeloop']\n",
    "    stats_file = os.path.join(base_dir, timeloop_dir, 'baseline_evaluation', sub_dir, \"layer{}\".format(layer_id_for_timeloop), \\\n",
    "                              \"timeloop-model.stats.txt\")\n",
    "    with open(stats_file, 'r') as f:\n",
    "        lines = f.read().split('\\n')[-200:]\n",
    "        for line in lines:\n",
    "            if line.startswith('Energy'):\n",
    "                energy = eval(line.split(': ')[1].split(' ')[0]) * float(10**6) # micro to pico\n",
    "            elif line.startswith('Cycles'):\n",
    "                cycle = eval(line.split(': ')[1])\n",
    "                \n",
    "    baseline_latency += cycle\n",
    "    baseline_energy += energy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "541f492c",
   "metadata": {},
   "source": [
    "#### 6-2. Calculate the latency and energy for secure accelerators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b96071d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_total_stats(cost_dict, rehash_cost_dict, n_layers, exclude_layers):\n",
    "    total_latency = 0\n",
    "    total_energy = 0\n",
    "    additional_off_chip_traffic = [0, 0, 0] # redundant, hash, rehash\n",
    "    \n",
    "    for idx in range(1, n_layers + 1):\n",
    "        if idx in layers_exclude_from_search:\n",
    "            continue\n",
    "        total_latency += cost_dict[idx]['total_latency']\n",
    "        total_energy += cost_dict[idx]['total_energy']\n",
    "        additional_off_chip_traffic[0] += cost_dict[idx]['total_redundant_bits']\n",
    "        additional_off_chip_traffic[1] += cost_dict[idx]['total_hash_bits']\n",
    "    \n",
    "    for key in rehash_cost_dict.keys():\n",
    "        idx1 = key[0]\n",
    "        idx2 = key[1]\n",
    "        \n",
    "        if idx1 in layers_exclude_from_search or idx2 in layers_exclude_from_search:\n",
    "            continue\n",
    "            \n",
    "        total_latency += rehash_cost_dict[key]['total_latency']\n",
    "        total_energy += rehash_cost_dict[key]['total_energy']\n",
    "        additional_off_chip_traffic[2] += rehash_cost_dict[key]['add_memory_traffic']\n",
    "        \n",
    "    return total_latency, total_energy, additional_off_chip_traffic\n",
    "\n",
    "cts_latency, cts_energy, cts_traffic = calculate_total_stats(cts_cost_dict, cts_rehash_cost_dict, n_layers, layers_exclude_from_search)\n",
    "cos_latency, cos_energy, cos_traffic = calculate_total_stats(cos_cost_dict, cos_rehash_cost_dict, n_layers, layers_exclude_from_search)\n",
    "coc_latency, coc_energy, coc_traffic = calculate_total_stats(coc_cost_dict, coc_rehash_cost_dict, n_layers, layers_exclude_from_search)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8b816f4",
   "metadata": {},
   "source": [
    "#### 6-3. Draw the figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33b8f555",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "\n",
    "def plot_bar_graph(ax, data, labels, show_legend, title, show_ylabel):\n",
    "    x_locations = np.arange(len(data))\n",
    "\n",
    "    # Set the width of each bar\n",
    "    num_groups = data.shape[0]\n",
    "    num_bars = data.shape[1]\n",
    "    bar_width = 1.0 / (num_bars + 3)\n",
    "    space_width = 0.05\n",
    "    \n",
    "    # Set the x-axis positions for each group\n",
    "    x_positions = np.arange(num_groups)\n",
    "\n",
    "    ax.set_prop_cycle('color', plt.cm.bone(np.linspace(0, 1, num_bars)))\n",
    "    \n",
    "    # Define the hatch patterns to use for each bar\n",
    "    hatch_patterns = ['.', '/', '\\\\', 'x', '-', '+']\n",
    "\n",
    "    # Plot each group of bars\n",
    "    for i in range(num_bars):\n",
    "        # Calculate the x-axis positions for each bar within each group\n",
    "        x_pos = x_positions + i * (bar_width + space_width)\n",
    "        rects = ax.bar(x_pos, data[:, i], width=bar_width, align='edge', label=labels[i], edgecolor='black')\n",
    "\n",
    "        # Add value of each bar as text\n",
    "        for rect in rects:\n",
    "            height = rect.get_height()\n",
    "            ax.annotate(f'{height:.2f}', xy=(rect.get_x() + rect.get_width() / 2, height),\n",
    "                        xytext=(0, 3),  # 3 points vertical offset\n",
    "                        textcoords=\"offset points\", ha='center', va='bottom', fontsize=14)\n",
    "\n",
    "    # Remove the y-axis ticks\n",
    "    ax.tick_params(axis='x', which='both', length=0, labelbottom=False)\n",
    "\n",
    "    # Set the y-axis label\n",
    "    if show_ylabel:\n",
    "        ax.set_ylabel('Normalized Latency', fontsize=15)\n",
    "    \n",
    "    ax.yaxis.grid(True, linestyle='--', which='major', color='grey', alpha=0.5)\n",
    "    \n",
    "    # Set the ylim\n",
    "    max_val = data.max() + max((data.max() - data.min()) * 0.4, 1)\n",
    "    ax.set_ylim(0, max_val)\n",
    "\n",
    "    if show_legend:\n",
    "        ax.legend(ncol=3, loc='lower center', fontsize=15, bbox_to_anchor=(0.5, -0.25))\n",
    "        \n",
    "    ax.set_title(title, fontsize=15, fontweight='bold', y=1.02)\n",
    "    \n",
    "fig, ax = plt.subplots(figsize=(4, 3))\n",
    "plot_bar_graph(ax, np.asarray([float(cts_latency) / float(baseline_latency), \\\n",
    "                               float(cos_latency) / float(baseline_latency), \\\n",
    "                               float(coc_latency) / float(baseline_latency)]).reshape((1, 3)), \\\n",
    "               ['Crypt-Tile-Single', 'Crypt-Opt-Single', 'Crypt-Opt-Cross'], \n",
    "               True, model_name, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fe82c99",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_stacked_bar_graph(ax, data, labels, show_legend, title, ylog, show_ylabel):\n",
    "    # Define the x locations for each group of bars\n",
    "    x_locations = np.arange(len(data))\n",
    "\n",
    "    # Set the width of each bar\n",
    "    num_groups = data[0].shape[0]\n",
    "    num_bars = data[0].shape[1]\n",
    "    bar_width = 1.0 / (num_bars + 3)\n",
    "    space_width = 0.05\n",
    "    \n",
    "    # Set the x-axis positions for each group\n",
    "    x_positions = np.arange(num_groups)\n",
    "    \n",
    "    hatches = ['/', '\\\\', '|']\n",
    "    cmap = plt.cm.get_cmap('bone', 4)\n",
    "    \n",
    "    # Plot each group of bars\n",
    "    for i in range(num_bars):\n",
    "        # Calculate the x-axis positions for each bar within each group\n",
    "        x_pos = x_positions + i * (bar_width + space_width)\n",
    "        bottom = np.zeros(num_bars)\n",
    "        for j in range(len(data)):\n",
    "            data_ = data[j]\n",
    "            \n",
    "            edgecolor = 'white' if i < 2 else 'black'\n",
    "            ax.bar(x_pos, data_[:, i], width=bar_width, align='edge', edgecolor=edgecolor, hatch=hatches[j], \\\n",
    "                   bottom=bottom, color=cmap(i))\n",
    "            ax.bar(x_pos, data_[:, i], width=bar_width, align='edge', edgecolor='black', \\\n",
    "                   bottom=bottom, color='none')\n",
    "            bottom += data_[:, i]\n",
    "\n",
    "    # Remove the y-axis ticks\n",
    "    ax.tick_params(axis='x', which='both', length=0, labelbottom=False)\n",
    "\n",
    "    # Set the ylim\n",
    "    ax.yaxis.grid(True, linestyle='--', which='major', color='grey', alpha=0.5)\n",
    "    \n",
    "    if ylog:\n",
    "        ax.set_yscale('log')\n",
    "        \n",
    "    if show_ylabel:\n",
    "        ax.set_ylabel('Additional \\n Off-chip Traffic (bits)', fontsize=15)\n",
    "        \n",
    "    ax.set_title(title, fontsize=15, fontweight='bold', y=1.02)\n",
    "\n",
    "    if show_legend:\n",
    "        # Add the legend\n",
    "        legend_elements = [plt.Rectangle((0, 0), 1, 1, facecolor='none', edgecolor='black', linewidth=1, hatch=hatches[i]) \\\n",
    "                           for i in range(len(data))]\n",
    "        # legend_handles = [Patch(facecolor='white', edgecolor='white', hatch=hatches[i]) for i in range(len(hatches))]\n",
    "        ax.legend(legend_elements, labels, \\\n",
    "                  ncol=3, bbox_to_anchor=(0.5, -0.25), loc='lower center', fontsize=15)\n",
    "        \n",
    "fig, ax = plt.subplots(figsize=(4, 3))\n",
    "plot_stacked_bar_graph(ax, \\\n",
    "                       [np.asarray([cts_traffic[2], cos_traffic[2], coc_traffic[2]]).reshape((1, 3)), \\\n",
    "                        np.asarray([cts_traffic[0], cos_traffic[0], coc_traffic[0]]).reshape((1, 3)), \\\n",
    "                        np.asarray([cts_traffic[1], cos_traffic[1], coc_traffic[1]]).reshape((1, 3))], \\\n",
    "                        ['Rehash', 'Redundant', 'Hash'], True, model_name, False, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e045b1d1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}