Classify miss-inplaced tensors in logs. (pytorch#139240)

Summary:
use signpost logs,
a followup is to remove the field possibly_missed_reinplacing_opportunities form dynamo compile table.

Differential Revision: D65180194

Pull Request resolved: pytorch#139240
Approved by: https://github.com/zou3519

Author: laithsakka

test/inductor/test_inplacing_pass.py

@@ -6,7 +6,7 @@
 import torch._inductor.config as inductor_config
 from functorch import make_fx
 from torch import Tensor
-from torch._dynamo.utils import counters
+from torch._dynamo.utils import ReinplaceCounters
 from torch._higher_order_ops.auto_functionalize import (
     auto_functionalized,
     auto_functionalized_v2,
@@ -31,11 +31,11 @@
 
 
 def num_reinplacing_failures():
-    return counters["inductor"]["possibly_missed_reinplacing_opportunities"]
+    return ReinplaceCounters.get_total_missed()
 
 
 def miss_inplaced_bytes():
-    return counters["inductor"]["possibly_missed_reinplacing_bytes"]
+    return ReinplaceCounters.get_total_missed_bytes()
 
 
 @torch.library.custom_op("_reinplacing::sin", mutates_args={"result"})
@@ -85,7 +85,7 @@ def boo(x: torch.Tensor) -> None:
 
 class TestReinplacingPassCorrectness(InductorTestCase):
     def setUp(self):
-        counters.clear()
+        ReinplaceCounters.clear()
         return super().setUp()
 
     def _test(self, f):
@@ -138,7 +138,7 @@ def f(x, y):
         self._test(f)
 
     def test_counters_functionalize_old(self):
-        counters.clear()
+        ReinplaceCounters.clear()
 
         def f(x):
             out = torch.empty_like(x)
@@ -158,7 +158,7 @@ def f(x):
         self.assertEqual(miss_inplaced_bytes(), 12)
 
     def test_counters_functionalize_v2(self):
-        counters.clear()
+        ReinplaceCounters.clear()
 
         def f(x):
             out = torch.empty_like(x)
@@ -314,7 +314,7 @@ def test_multi_output_intermediate(self):
                 with inductor_config.patch(
                     {"enable_auto_functionalized_v2": enable_v2}
                 ):
-                    counters.clear()
+                    ReinplaceCounters.clear()
 
                     def f(x):
                         out1 = torch.empty_like(x)
@@ -329,7 +329,7 @@ def f(x):
                     self.assertEqual(num_reinplacing_failures(), 0)
 
     def test_multiple_mutations(self):
-        counters.clear()
+        ReinplaceCounters.clear()
 
         def f(x, out):
             sin(x, out)
@@ -345,7 +345,7 @@ def f(x, out):
         self.assertEqual(num_reinplacing_failures(), 0)
 
     def test_multiple_intermediate(self):
-        counters.clear()
+        ReinplaceCounters.clear()
 
         def f(x):
             out = torch.empty_like(x)

torch/_dynamo/convert_frame.py

@@ -971,12 +971,7 @@ def format_guard_failures() -> str:
         fail_reason: Optional[str] = None
         fail_user_frame_filename: Optional[str] = None
         fail_user_frame_lineno: Optional[int] = None
-        start_possibly_missed_reinplacing_opportunities = torch._dynamo.utils.counters[
-            "inductor"
-        ]["possibly_missed_reinplacing_opportunities"]
-        start_possibly_missed_reinplacing_bytes = torch._dynamo.utils.counters[
-            "inductor"
-        ]["start_possibly_missed_reinplacing_bytes"]
+        torch._dynamo.utils.ReinplaceCounters.clear()
         guarded_code = None
         try:
             guarded_code = compile_inner(code, one_graph, hooks, transform)
@@ -1054,33 +1049,17 @@ def format_guard_failures() -> str:
                 compliant_custom_ops = {
                     op.__qualname__ for op in output.compliant_custom_ops
                 }
-                possibly_missed_reinplacing_opportunities = (
-                    torch._dynamo.utils.counters["inductor"][
-                        "possibly_missed_reinplacing_opportunities"
-                    ]
-                    - start_possibly_missed_reinplacing_opportunities
-                )
                 remote_cache_time_saved = frame_phase_timing[frame_key].get(
                     "remote_cache_time_saved", 0
                 )
-                possibly_missed_reinplacing_bytes = (
-                    torch._dynamo.utils.counters["inductor"][
-                        "possibly_missed_reinplacing_bytes"
-                    ]
-                    - start_possibly_missed_reinplacing_bytes
-                )
-                if possibly_missed_reinplacing_bytes != 0:
-                    signpost_event(
-                        "inductor",
-                        "auto_functionalize",
-                        {"missed_reinplacing_bytes": possibly_missed_reinplacing_bytes},
-                    )
                 remote_fx_graph_cache_get_time = frame_phase_timing[frame_key].get(
                     "remote_fx_graph_cache_get", None
                 )
                 remote_fx_graph_cache_put_time = frame_phase_timing[frame_key].get(
                     "remote_fx_graph_cache_put", None
                 )
+                torch._dynamo.utils.ReinplaceCounters.log()
+
             else:
                 guard_count = None
                 shape_env_guard_count = None
@@ -1096,7 +1075,6 @@ def format_guard_failures() -> str:
                 restart_reasons = set()
                 # If compilation failed, the entire time is wasted
                 dynamo_time_before_restart = duration_ns / 1e9
-                possibly_missed_reinplacing_opportunities = None
                 remote_cache_time_saved = None
                 remote_fx_graph_cache_get_time = None
                 remote_fx_graph_cache_put_time = None
@@ -1161,7 +1139,6 @@ def clean_for_json(d: Dict[str, Any]) -> Dict[str, Any]:
                 restart_reasons,
                 dynamo_time_before_restart,
                 guarded_code is not None,
-                possibly_missed_reinplacing_opportunities,
                 remote_cache_time_saved,
                 structured_logging_overhead_s,
                 config.suppress_errors,

torch/_dynamo/utils.py

@@ -73,7 +73,11 @@
 from torch._dispatch.python import enable_python_dispatcher
 from torch._guards import Source, TracingContext
 from torch._subclasses.meta_utils import is_sparse_compressed
-from torch._utils_internal import log_chromium_event_internal, log_compilation_event
+from torch._utils_internal import (
+    log_chromium_event_internal,
+    log_compilation_event,
+    signpost_event,
+)
 from torch.fx._utils import _format_graph_code, lazy_format_graph_code
 from torch.nn.modules.lazy import LazyModuleMixin
 from torch.utils._triton import has_triton, has_triton_package
@@ -143,6 +147,51 @@
 timer_counter = itertools.count()
 
 
+# Abstraction on top of counters.
+class ReInplaceTrigger(enum.Enum):
+    AUTO_FUNC_V1 = 1
+    AUTO_FUNC_V2 = 2
+    TRITON_OPS = 3
+
+
+class ReinplaceCounters:
+    _values: DefaultDict[str, int] = collections.defaultdict(int)
+
+    # Track sizes of known not re-inplaced tensors (exclude dynamic shapes).
+    @classmethod
+    def add_missed_bytes(cls, trigger: ReInplaceTrigger, bytes: int):
+        cls._values[f"missed_bytes_{trigger.name}"] += bytes
+
+    # Track number of not re-inplaced tensors.
+    @classmethod
+    def add_missed_opportunities(cls, trigger: ReInplaceTrigger, count: int):
+        cls._values[f"missed_tensors_{trigger}"] += count
+
+    @classmethod
+    def clear(cls):
+        cls._values.clear()
+
+    @classmethod
+    def get_total_missed(cls):
+        sum = 0
+        for trigger in ReInplaceTrigger:
+            sum += cls._values.get(f"missed_tensors_{trigger}", 0)
+        return sum
+
+    @classmethod
+    def get_total_missed_bytes(cls):
+        sum = 0
+        for trigger in ReInplaceTrigger:
+            sum += cls._values.get(f"missed_bytes_{trigger.name}", 0)
+        return sum
+
+    @classmethod
+    def log(cls):
+        # if not empty log.
+        if cls._values:
+            signpost_event("inductor", "reinplace_counters", cls._values)
+
+
 def tabulate(
     rows: Union[List[Tuple[str, object]], List[List[object]]],
     headers: Union[Tuple[str, ...], List[str]],
@@ -843,7 +892,6 @@ class CompilationMetrics:
     # to install any guarded code.  True means we actually decided to install
     # a compiled frame
     has_guarded_code: Optional[bool] = None
-    possibly_missed_reinplacing_opportunities: Optional[int] = None
     remote_cache_time_saved_s: Optional[float] = None
     structured_logging_overhead_s: Optional[float] = None
     config_suppress_errors: Optional[bool] = None

torch/_inductor/fx_passes/reinplace.py

@@ -8,6 +8,7 @@
 
 import torch
 from torch._dispatch.python import enable_python_dispatcher
+from torch._dynamo.utils import ReinplaceCounters, ReInplaceTrigger
 from torch._higher_order_ops.triton_kernel_wrap import (
     kernel_side_table,
     triton_kernel_wrapper_functional,
@@ -497,7 +498,12 @@ def can_inplace(node, mutated_arg):
             )
 
     def log_inplace_results(
-        node_name, old_tensors_to_clone, tensors_to_clone, missed_args, missed_nodes
+        node_name,
+        old_tensors_to_clone,
+        tensors_to_clone,
+        missed_args,
+        missed_nodes,
+        trigger,
     ):
         # Total size of possibly_missed_reinplacing_opportunities for tensors with static shapes.
         missed_bytes = 0
@@ -531,17 +537,14 @@ def bytes(node):
             missed_args,
             missed_bytes,
         )
-        torch._dynamo.utils.counters["inductor"][
-            "possibly_missed_reinplacing_opportunities"
-        ] += len(missed_args)
-        torch._dynamo.utils.counters["inductor"][
-            "possibly_missed_reinplacing_bytes"
-        ] += missed_bytes
+
+        ReinplaceCounters.add_missed_opportunities(trigger, len(missed_args))
+        ReinplaceCounters.add_missed_bytes(trigger, missed_bytes)
 
     replace_dict: Dict[torch.fx.Node, torch.fx.Node] = {}
 
     def reinplace_and_refine_tensors_to_clone(
-        old_tensors_to_clone, kwargs, node_name, auto_functionalize_v2=False
+        old_tensors_to_clone, kwargs, node_name, trigger
     ):
         tensors_to_clone: List[str] = []
         storage_of_reinplaced_args = set()
@@ -580,7 +583,7 @@ def tensor_with_same_storage_already_reinplaced(arg):
                 copy_node = copy_args_to_copy_nodes.get((mutated_arg, node))
                 if copy_node is not None:
                     replace_dict[copy_node] = copy_node.args[0]
-                if not auto_functionalize_v2:
+                if not trigger == ReInplaceTrigger.AUTO_FUNC_V2:
                     for user in node.users:
                         # For auto_functionalize_v2, arg is the index of the base, where base at index i corresponds to
                         # output atindex size(out)+i.
@@ -602,7 +605,12 @@ def tensor_with_same_storage_already_reinplaced(arg):
                 tensors_to_clone.append(arg)
 
         log_inplace_results(
-            node_name, old_tensors_to_clone, tensors_to_clone, missed_args, missed_nodes
+            node_name,
+            old_tensors_to_clone,
+            tensors_to_clone,
+            missed_args,
+            missed_nodes,
+            trigger,
         )
         return tensors_to_clone
 
@@ -628,7 +636,7 @@ def tensor_with_same_storage_already_reinplaced(arg):
                 bases_to_clone,
                 base_tensors_dct,
                 node.target,
-                auto_functionalize_v2=True,
+                ReInplaceTrigger.AUTO_FUNC_V2,
             )
             # Stash the metadata. There is a pass later on where we decompose
             # auto_functionalized into clones + a mutable op; this metadata
@@ -647,7 +655,7 @@ def tensor_with_same_storage_already_reinplaced(arg):
                 tensors_to_clone,
                 node.kwargs,
                 _mutable_op._name,
-                auto_functionalize_v2=False,
+                ReInplaceTrigger.AUTO_FUNC_V1,
             )
 
             # Stash the metadata. There is a pass later on where we decompose
@@ -679,7 +687,10 @@ def tensor_with_same_storage_already_reinplaced(arg):
             # This pass iterates over them and sees which ones are safe
             # to eliminate (i.e. no longer need the clones)
             tensors_to_clone = reinplace_and_refine_tensors_to_clone(
-                node.kwargs["tensors_to_clone"], node.kwargs["kwargs"], kernel_name
+                node.kwargs["tensors_to_clone"],
+                node.kwargs["kwargs"],
+                kernel_name,
+                ReInplaceTrigger.TRITON_OPS,
             )
 
             kwargs = dict(node.kwargs)