Improve cache key graph printing performance (pytorch#151928)

Teach the graph printer how to allow overriding printing SymTypes (`SymInt`, `SymFloat`, `SymBool`) and then use that to reuse the fast SymNode printing from `torch._inductor.utils.sympy_str()` to make computing the cache key faster.

On my computer the repro from pytorch#151823 goes from 480s -> 80s (still terrible... but better).

Fixes pytorch#151823

Pull Request resolved: pytorch#151928
Approved by: https://github.com/laithsakka

Author: aorenste

torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py

@@ -186,7 +186,10 @@ def aot_dispatch_base(
     aot_forward_graph_str = None
     if aot_config.cache_info is not None:
         aot_forward_graph_str = fw_module.print_readable(
-            print_output=False, include_stride=True, include_device=True
+            print_output=False,
+            include_stride=True,
+            include_device=True,
+            fast_sympy_print=True,
         )
 
     fakified_out_wrapper = FakifiedOutWrapper()

torch/_inductor/compile_fx.py

@@ -844,9 +844,11 @@ def _compile_fx_inner(
             assert mb_compiled_graph is None
             log.debug(
                 "FX cache bypass reason: %s",
-                cache_info.get("cache_bypass_reason", "unknown")
-                if cache_info is not None
-                else "FX cache disabled or key generation failed",
+                (
+                    cache_info.get("cache_bypass_reason", "unknown")
+                    if cache_info is not None
+                    else "FX cache disabled or key generation failed"
+                ),
             )
             mb_compiled_graph = fx_codegen_and_compile(
                 gm, example_inputs, inputs_to_check, **graph_kwargs
@@ -1167,8 +1169,15 @@ def codegen_and_compile(
                         colored=True,
                     ),
                 )
+
+                # We're printing the graph to be used as a cache key - so a
+                # printer which is a little less readable but faster is
+                # appropriate.
                 inductor_post_grad_graph_str = gm.print_readable(
-                    print_output=False, include_stride=True, include_device=True
+                    print_output=False,
+                    include_stride=True,
+                    include_device=True,
+                    fast_sympy_print=True,
                 )
                 trace_structured(
                     "inductor_post_grad_graph",
@@ -1268,12 +1277,12 @@ def codegen_and_compile(
                     is_inference=is_inference,
                     is_backward=is_backward,
                     const_output_index=const_output_index,
-                    const_wrapper_code=const_wrapper_code.value
-                    if const_wrapper_code
-                    else None,
-                    const_kernel_code=const_kernel_code.value
-                    if const_kernel_code
-                    else None,
+                    const_wrapper_code=(
+                        const_wrapper_code.value if const_wrapper_code else None
+                    ),
+                    const_kernel_code=(
+                        const_kernel_code.value if const_kernel_code else None
+                    ),
                     const_module=const_graph,
                     inputs_to_check=inputs_to_check,
                 )

torch/fx/graph.py

@@ -12,7 +12,7 @@
 import typing
 import warnings
 from collections import defaultdict
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator
 from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Any, Callable, Literal, NamedTuple, Optional, TYPE_CHECKING
@@ -317,6 +317,9 @@ def _parse_stack_trace(stack_trace: str):
 
 @compatibility(is_backward_compatible=False)
 class CodeGen:
+    # This is an override hook so we can customize the SymNode printer.
+    _sym_repr: Callable[["torch.types.PySymType"], str] = lambda x: repr(x)
+
     def __init__(self):
         self._body_transformer: Optional[TransformCodeFunc] = None
         self._func_name: str = "forward"
@@ -609,7 +612,8 @@ def emit_node(node: Node):
                         f'{dim_blue(stride_annotation)}{dim_green(device_annotation)}"'
                     )
                 elif isinstance(meta_val, py_sym_types):
-                    maybe_type_annotation = f': "Sym({meta_val})"'
+                    val_str = CodeGen._sym_repr(meta_val)
+                    maybe_type_annotation = f': "Sym({val_str})"'
                 elif isinstance(meta_val, TensorMetadata):
                     maybe_type_annotation = f': "{dtype_abbrs[meta_val.dtype]}{stringify_shape(meta_val.shape)}"'
 
@@ -1907,6 +1911,18 @@ def on_generate_code_context_manager():
         return on_generate_code_context_manager()
 
 
+@contextmanager
+def _override_sym_repr(
+    override: Callable[["torch.types.PySymType"], str]
+) -> Iterator[None]:
+    tmp = CodeGen._sym_repr
+    try:
+        CodeGen._sym_repr = override
+        yield
+    finally:
+        CodeGen._sym_repr = tmp
+
+
 def _identity(x):
     return x
 

torch/fx/graph_module.py

@@ -17,7 +17,14 @@
 from torch.package import Importer, PackageExporter, PackageImporter, sys_importer
 
 from ._compatibility import compatibility
-from .graph import _custom_builtins, _is_from_torch, _PyTreeCodeGen, Graph, PythonCode
+from .graph import (
+    _custom_builtins,
+    _is_from_torch,
+    _override_sym_repr,
+    _PyTreeCodeGen,
+    Graph,
+    PythonCode,
+)
 
 
 __all__ = [
@@ -927,18 +934,33 @@ def print_readable(
         include_stride=False,
         include_device=False,
         colored=False,
+        *,
+        # If `fast_sympy_print` is True then we use a sympy printer which is faster
+        # but may result in less-readable output.
+        fast_sympy_print: bool = False,
     ):
         """
         Return the Python code generated for current GraphModule and its children GraphModules
         """
-        return _print_readable(
-            self,
-            self._get_name(),
-            print_output,
-            include_stride,
-            include_device,
-            colored,
-        )
+        ctx_mgr = contextlib.ExitStack()
+        with ctx_mgr:
+            if fast_sympy_print:
+                from torch._inductor.utils import sympy_str
+
+                def fast_repr(expr: torch.types.PySymType) -> str:
+                    return sympy_str(expr.node.expr)
+
+                ctx_mgr.enter_context(_override_sym_repr(fast_repr))
+
+            r = _print_readable(
+                self,
+                self._get_name(),
+                print_output,
+                include_stride,
+                include_device,
+                colored,
+            )
+            return r
 
     def __str__(self) -> str:
         orig_str = super().__str__()