.

Author: hylin2002

.vscode/settings.json

@@ -0,0 +1,4 @@
+{
+    "python-envs.defaultEnvManager": "ms-python.python:pyenv",
+    "python-envs.pythonProjects": []
+}

src/benchmark_utils.py

@@ -18,13 +18,12 @@
 import subprocess
 import shutil
 
-# The dictionary to map a CPU collective function to its corresponding operation on TPU
-# "psum_scatter_ici_op" has different implementation according to its `matrix_dim` and the number of TPUs, so it's not considered in this mapping dictionary.
+# The dictionary to map a JAX (collective) function to its main HLO.
 TARGET_TASK_NAME_COLLECTIVES_MAP = {
     "all_to_all_ici_op": r"all-to-all.[0-9]+",
     "all_gather_ici_op": r"all-gather.[0-9]+",
     "psum_ici_op": r"all-reduce.[0-9]+",
-    "ppermute_ici_op": r"collective-permute-done",
+    "ppermute_ici_op": r"collective-permute.[0-9]+",
 }
 
 def simple_timeit(f, *args, matrix_dim=None, tries=10, task=None, trace_dir=None) -> float:
@@ -72,7 +71,7 @@ def get_metrics_from_trace(trace: dict[str, Any], task: str) -> list[float]:
 
     # Check if the given task name is a collective with corresponding TPU opertion.
     # This is a workaround and should be reverted or refactored in future.
-    if task in TARGET_TASK_NAME_COLLECTIVES_MAP.keys():
+    if task in TARGET_TASK_NAME_COLLECTIVES_MAP:
         task = TARGET_TASK_NAME_COLLECTIVES_MAP[task]
         return get_metrics_from_trace_tpu(trace, task)
     event_matcher = re.compile(task)