Do not clear UncachedCompile between graph executions (#9282)

Co-authored-by: Haifeng Jin <[email protected]>

Author: haifeng-jin

test/dynamo/test_bridge.py

@@ -303,6 +303,20 @@ def foo(t):
     t = torch.randint(0, 5, (3,), device=device)
     self._compile_and_check(foo, (t,))
 
+  def test_metrics_preserved(self):
+    metrics.clear_counters()
+    torch_xla._XLAC._xla_increment_counter('UncachedCompile', 3)
+    torch_xla._XLAC._xla_increment_counter('DynamoExtractCompiledGraph', 4)
+
+    model = BasicModule()
+    graph_module = torch.fx.symbolic_trace(model)
+    collector = bridge.UnsupportedNodesCollector(graph_module)
+    collector.run(torch.randn(1), torch.randn(1))
+
+    self.assertEqual(metrics.counter_value('UncachedCompile'), 3)
+    self.assertEqual(metrics.counter_value('DynamoExtractCompiledGraph'), 4)
+    metrics.clear_counters()
+
 
 if __name__ == "__main__":
   from torch._dynamo.test_case import run_tests

torch_xla/_dynamo/dynamo_bridge.py

@@ -646,6 +646,7 @@ def run_node(self, n: torch.fx.Node):
     # We need to restore this metric count later, so save it in a separate variable
     dynamo_extract_graph_helper_metric_count = metrics.counter_value(
         'DynamoExtractCompiledGraph')
+    uncached_compile_metric_count = metrics.counter_value('UncachedCompile')
 
     metrics.clear_counters()
     result = super().run_node(n)
@@ -684,6 +685,9 @@ def all_tensors_on_xla_device(value):
     # Restore this metric counter
     torch_xla._XLAC._xla_increment_counter(
         'DynamoExtractCompiledGraph', dynamo_extract_graph_helper_metric_count)
+    if uncached_compile_metric_count is not None:
+      torch_xla._XLAC._xla_increment_counter('UncachedCompile',
+                                             uncached_compile_metric_count)
 
     return result