From 51504cc0bbec11696190528cde1404f6e53a452b Mon Sep 17 00:00:00 2001 From: Janani Sriram Date: Tue, 11 Jun 2024 20:30:00 -0700 Subject: [PATCH] Add support for x_only = True extra metrics during table generation (#2294) Summary: Pull Request resolved: https://github.com/pytorch/benchmark/pull/2294 Add support for extra metrics that are the same across all benchmark implementations, identified by the `x_only = True` argument for [`register_metric()`](https://www.internalfb.com/code/fbsource/[dde2970b470a85612b904a14825aa7f0998f1658]/fbcode/pytorch/benchmark/torchbenchmark/util/triton_op.py?lines=328). This diff supports additional functionality in D58396957, where the current implementation did not support adding new metrics with `x_only = True`. Reviewed By: xuzhao9 Differential Revision: D58435001 --- torchbenchmark/util/triton_op.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/torchbenchmark/util/triton_op.py b/torchbenchmark/util/triton_op.py index 07da2198ad..6d763b6ca3 100644 --- a/torchbenchmark/util/triton_op.py +++ b/torchbenchmark/util/triton_op.py @@ -210,7 +210,10 @@ def select_metric(m): # Append x_val_only metrics for x_only_metric in x_only_metrics: x_only_metric_dict = asdict(y_val[y_val_keys[0]]) - row.append(x_only_metric_dict[x_only_metric]) + 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]) for k, _label in y_val_keys: metrics_dict = asdict(y_val[k]) if metrics_dict["error_msg"]: