Feat/replay from report (#168)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 2 people

TraceLens/EventReplay/batched_replay.py

@@ -3,7 +3,7 @@
 import argparse
 import warnings
 import torch
-from TraceLens.EventReplay.utils import TensorCfg, build_tensor, benchmark_func, summarize_tensor, dict_profile2torchdtype
+from .utils import TensorCfg, build_tensor, benchmark_func, summarize_tensor, dict_profile2torchdtype
 
 def _get_args_kwargs_from_ir(event_replay_IR: dict[str, any], device: str = 'cuda') -> tuple[list[any], dict[str, any]]:
     # (Copy the implementation of _get_args_kwargs_from_ir from Step 1 here)

examples/event_replayer_example.ipynb

@@ -20,6 +20,7 @@
    "outputs": [],
    "source": [
     "# replace by your profile path, it can be a single rank profile from a multi gpu run as well\n",
+    "# If you are interested in replaying from perf report then skip down \n",
     "path =\"/home/buffer/resnet_trace.json\"\n",
     "perf_analyzer = TreePerfAnalyzer.from_file(path)"
    ]
@@ -32,8 +33,7 @@
    "outputs": [],
    "source": [
     "# Replay works for any op, taking gemm as example \n",
-    "# gemm_events = [event for event in perf_analyzer.tree.events if event['name'] in ['aten::addmm', 'aten::mm', 'aten::_scaled_mm']]\n",
-    "gemm_events = [event for event in perf_analyzer.tree.events if event['name'] in ['aten::convolution']]\n",
+    "gemm_events = [event for event in perf_analyzer.tree.events if event['name'] in ['aten::addmm', 'aten::mm', 'aten::_scaled_mm']]\n",
     "df_gemm_ops = perf_analyzer.build_df_perf_metrics(gemm_events)\n",
     "df_gemm_summary = perf_analyzer.summarize_df_perf_metrics(df_gemm_ops, ['mean'])\n",
     "df_gemm_summary"
@@ -218,14 +218,98 @@
     "print(f\"Created zip file: {zip_file_path}\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "1722250f",
+   "metadata": {},
+   "source": [
+    "Replay ops from report"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "144ce15b-5f0a-4f8e-b7a2-f287154a5d18",
    "metadata": {},
    "outputs": [],
    "source": [
-    "batched_replay.__file__"
+    "# we can replay events from the perf reports as well - without the full profile too!\n",
+    "# This is because we essentially require the args and the op name to replay\n",
+    "# excel -> df -> for each row (row -> event -> replayer -> replayer IR -> append to replayer IR list) -> save replayer IR list as json\n",
+    "import pandas as pd\n",
+    "import ast\n",
+    "# read sheet from excel\n",
+    "\n",
+    "df_unique_ops = pd.read_excel('/path/to/your/perf_report.xlsx', sheet_name='sheet_name')\n",
+    "\n",
+    "def row_to_evt(row):\n",
+    "    event = {\n",
+    "        'name': row['name'],\n",
+    "        'args': {\n",
+    "            'Input Dims': ast.literal_eval(row['Input Dims']),\n",
+    "            'Input Strides': ast.literal_eval(row['Input Strides']),\n",
+    "            'Input type': ast.literal_eval(row['Input type']),\n",
+    "            'Concrete Inputs': ast.literal_eval(row['Concrete Inputs']),\n",
+    "        }\n",
+    "    }\n",
+    "    return event\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eb91e0c0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "repro_data_list = []\n",
+    "processed_count = 0\n",
+    "# lets say we are interested in the following ops\n",
+    "ops_interest = ['aten::miopen_convolution',\n",
+    "                'aten::convolution_backward', \n",
+    "                'aten::miopen_batch_norm',\n",
+    "                'aten::miopen_batch_norm_backward'] \n",
+    "\n",
+    "df_ops_interest = df_unique_ops[df_unique_ops['name'].isin(ops_interest)].copy()\n",
+    "\n",
+    "for index, row in df_ops_interest.iterrows():\n",
+    "    event = row_to_evt(row)\n",
+    "    # Initialize EventReplayer similar to above\n",
+    "    replayer = EventReplayer(event, lazy=True, verbose=False)\n",
+    "    # Extract the serializable info\n",
+    "    repro_info = replayer.get_repro_info()\n",
+    "    repro_data_list.append(repro_info)\n",
+    "    processed_count += 1\n",
+    "print(f\"Processed {processed_count} events.\")\n",
+    "# --- Save the Extracted Data ---\n",
+    "OUTPUT_REPRO_FILE = '/path/to/your/output/replay_ir.json'\n",
+    "if repro_data_list:\n",
+    "    print(f\"\\nSaving {len(repro_data_list)} extracted operator infos to '{OUTPUT_REPRO_FILE}'...\")\n",
+    "    with open(OUTPUT_REPRO_FILE, 'w') as f:\n",
+    "        json.dump(repro_data_list, f, indent=4)\n",
+    "    print(\"Save complete.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a0245187",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# STANDALONE ARTIFACTS FOR REPRO -  independent of model code or tracelens code\n",
+    "# artifacts include (a)replay_ir.json, (b) utils.py, (c) batched_replay.py\n",
+    "files = [\n",
+    "    OUTPUT_REPRO_FILE,\n",
+    "    tl_utils.__file__,  # Path to utils.py\n",
+    "    batched_replay.__file__,  # Path to batched_replay.py\n",
+    "    batched_replay.__file__.replace('batched_replay.py', 'batched_replay_readme.md') # path to the readme\n",
+    "]\n",
+    "zip_file_path = '/path/to/your/directory/replay_code.zip'  # Specify your desired zip file path\n",
+    "with zipfile.ZipFile(zip_file_path, 'w') as zipf:\n",
+    "    for file in files:\n",
+    "        zipf.write(file, arcname=os.path.basename(file))  # ← use file.name\n",
+    "print(f\"Created zip file: {zip_file_path}\")"
    ]
   }
  ],