[LTC] Upstream short_metrics (#89186)

Summary:
This pull request upstreams pytorch/xla#4148.

Test Plan:
xla CI.

Pull Request resolved: #89186
Approved by: https://github.com/JackCaoG

test/lazy/test_ts_opinfo.py

@@ -71,20 +71,28 @@ def init_lists():
         'linalg_pinv.atol_rtol_tensor',
         'logsumexp',
     ])
+    # For some ops, we don't support all variants. Here we use formatted_name
+    # to uniquely identify the variant.
+    SKIP_VARIANT_LIST = set([
+        'norm_nuc',
+        'min_reduction_with_dim'
+    ])
 
     return (LAZY_OPS_LIST,
             FALLBACK_LIST,
             SKIP_RUNTIME_ERROR_LIST,
             SKIP_INCORRECT_RESULTS_LIST,
             FUNCTIONAL_DECOMPOSE_LIST,
-            HAS_SYMINT_SUFFIX)
+            HAS_SYMINT_SUFFIX,
+            SKIP_VARIANT_LIST)
 
 (LAZY_OPS_LIST,
  FALLBACK_LIST,
  SKIP_RUNTIME_ERROR_LIST,
  SKIP_INCORRECT_RESULTS_LIST,
  FUNCTIONAL_DECOMPOSE_LIST,
- HAS_SYMINT_SUFFIX) = init_lists()
+ HAS_SYMINT_SUFFIX,
+ SKIP_VARIANT_LIST) = init_lists()
 
 torch.manual_seed(42)
 
@@ -166,6 +174,7 @@ class TestLazyOpInfo(TestCase):
           if op.name in LAZY_OPS_LIST
           and op.name not in SKIP_RUNTIME_ERROR_LIST
           and op.name not in FUNCTIONAL_DECOMPOSE_LIST
+          and op.formatted_name not in SKIP_VARIANT_LIST
           ], allowed_dtypes=(torch.float,))
     def test_dispatched_to_lazy(self, device, dtype, op):
         def get_name(op):

torch/csrc/lazy/core/metrics.cpp

@@ -172,7 +172,9 @@ std::vector<std::string> MetricsArena::GetCounterNames() {
   std::vector<std::string> names;
   std::lock_guard<std::mutex> lock(lock_);
   for (auto& name_data : counters_) {
-    names.push_back(name_data.first);
+    if (name_data.second->Value() > 0) {
+      names.push_back(name_data.first);
+    }
   }
   return names;
 }
@@ -353,6 +355,34 @@ std::string CreateMetricReport() {
   return ss.str();
 }
 
+std::string CreateMetricReport(
+    const std::vector<std::string>& counter_names,
+    const std::vector<std::string>& metric_names) {
+  MetricsArena* arena = MetricsArena::Get();
+  std::stringstream ss;
+  for (const std::string& metric_name : metric_names) {
+    MetricData* data = arena->GetMetric(metric_name);
+    if (data && data->TotalSamples() > 0) {
+      EmitMetricInfo(metric_name, data, &ss);
+    }
+  }
+  for (const std::string& counter_name : counter_names) {
+    CounterData* data = arena->GetCounter(counter_name);
+    if (data && data->Value() > 0) {
+      EmitCounterInfo(counter_name, data, &ss);
+    }
+  }
+  static std::string fall_back_counter_prefix = "aten::";
+  arena->ForEachCounter([&ss](const std::string& name, CounterData* data) {
+    if (name.rfind(fall_back_counter_prefix, 0) == 0 && data->Value() > 0) {
+      // it might emit duplicated counter if user also specified exact aten
+      // counter in the `counter_names` but it should be very rare.
+      EmitCounterInfo(name, data, &ss);
+    }
+  });
+  return ss.str();
+}
+
 std::vector<std::string> GetMetricNames() {
   return MetricsArena::Get()->GetMetricNames();
 }

torch/csrc/lazy/core/metrics.h

@@ -216,6 +216,11 @@ class TORCH_API Counter {
 // Creates a report with the current metrics statistics.
 TORCH_API std::string CreateMetricReport();
 
+// Creates a report with the selected metrics statistics.
+TORCH_API std::string CreateMetricReport(
+    const std::vector<std::string>& counter_names,
+    const std::vector<std::string>& metric_names);
+
 // Returns the currently registered metric names. Note that the list can grow
 // since metrics are usually function intialized (they are static function
 // variables).