Extend support for x_only = True extra metrics through canonical name…

… retrieval (#2300)

Summary:
Pull Request resolved: #2300

Index into the canonical name for a metric function, where function names are defined as tuples in the form `(canonical function name, customized function name)`. This diff extends support for extra metrics that are the same across all benchmark implementations, identified by the parameter `x_only = True`.

Reviewed By: xuzhao9

Differential Revision: D58552578

fbshipit-source-id: c9fc182de5a3c20d3555addf587ae543b7fbcb54

torchbenchmark/util/triton_op.py

@@ -209,8 +209,13 @@ def select_metric(m):
  row.append(x_val)
  # Append x_val_only metrics
  for x_only_metric in x_only_metrics:
- x_only_metric_dict = asdict(y_val[y_val_keys[0]])
- if "extra_metrics" in x_only_metric_dict and x_only_metric in x_only_metric_dict["extra_metrics"]:
+ x_only_metric_dict = asdict(
+ y_val[y_val_keys[0][0]]
+ ) # retrieve canonical name for metric function, where y_val_keys[0] = (canonical name, customized label name)
+ if (
+ "extra_metrics" in x_only_metric_dict
+ and x_only_metric in x_only_metric_dict["extra_metrics"]
+ ):
  row.append(x_only_metric_dict["extra_metrics"][x_only_metric])
  else:
  row.append(x_only_metric_dict[x_only_metric])