Add support for x_only = True extra metrics during table generation (#…

…2294)

Summary:
Pull Request resolved: #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

fbshipit-source-id: 887a66c08d0eef8c4c9c06f113d834f81ee572d9

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"]: