[export] Log private api uses. (#119848)

zhxchen17 · pytorchmergebot · commit 8f27fde2f54e · 2024-02-14T22:58:23.000Z
Summary: as title. The following APIs are logged: - capture_preautograd_graph - torch._export.aot_compile - external usage of _export_to_torch_ir (AOTInductor, Pippy) - constraints API - public use of torch._dynamo.export Test Plan: CI Differential Revision: D53735599 Pull Request resolved: #119848 Approved by: https://github.com/suo
diff --git a/torch/_dynamo/eval_frame.py b/torch/_dynamo/eval_frame.py
@@ -43,6 +43,7 @@
 import torch.utils.checkpoint
 from torch import _guards
 from torch._subclasses import fake_tensor
+from torch._utils_internal import log_export_usage
 from torch.export import Constraint
 from torch.export.dynamic_shapes import _process_dynamic_shapes
 from torch.fx.experimental.proxy_tensor import make_fx, maybe_disable_fake_tensor_mode
@@ -1115,6 +1116,7 @@ def export(
     assume_static_by_default: bool = False,
     same_signature: bool = True,
     disable_constraint_solver: bool = False,
+    _log_export_usage: bool = True,
     **extra_kwargs,
 ) -> Callable[..., ExportResult]:
     """
@@ -1179,6 +1181,9 @@ def export(
 
     Note - this headerdoc was authored by ChatGPT, with slight modifications by the author.
     """
+    if _log_export_usage:
+        log_export_usage(event="export.private_api", flags={"_dynamo"})
+
     # Deal with "local variable referenced before assignment"
     _f = f
     _assume_static_by_default = assume_static_by_default
diff --git a/torch/_export/__init__.py b/torch/_export/__init__.py
@@ -31,15 +31,25 @@
 from torch._functorch.aot_autograd import aot_export_module, GraphSignature
 from torch._functorch.eager_transforms import functionalize
 from torch._guards import detect_fake_mode
+from torch._inductor import config
 from torch._ops import OpOverload
 from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
 from torch._subclasses.functional_tensor import FunctionalTensor
+from torch._utils_internal import log_export_usage
 from torch.export._tree_utils import reorder_kwargs
+from torch.export._unlift import _create_stateful_graph_module
+from torch.export.dynamic_shapes import (
+    _process_constraints,
+    _process_dynamic_shapes,
+    Constraint,
+    dims,
+    dynamic_dim,
+)
 from torch.export.exported_program import (
+    _disable_prexisiting_fake_mode,
     ExportedProgram,
     ModuleCallEntry,
     ModuleCallSignature,
-    _disable_prexisiting_fake_mode,
 )
 from torch.export.graph_signature import (
     _sig_to_specs,
@@ -53,14 +63,6 @@
     SymIntArgument,
     TensorArgument,
 )
-from torch.export.dynamic_shapes import (
-    Constraint,
-    dims,
-    dynamic_dim,
-    _process_constraints,
-    _process_dynamic_shapes,
-)
-from torch.export._unlift import _create_stateful_graph_module
 from torch.fx import traceback as fx_traceback
 from torch.fx._compatibility import compatibility
 from torch.fx.experimental.proxy_tensor import make_fx, maybe_disable_fake_tensor_mode
@@ -77,7 +79,6 @@
     _AddRuntimeAssertionsForInlineConstraintsPass,
 )
 from .wrappers import _wrap_submodules
-from torch._inductor import config
 
 
 @dataclasses.dataclass
@@ -139,6 +140,8 @@ def capture_pre_autograd_graph(
     from torch.export._trace import _convert_input_to_fake, DEFAULT_EXPORT_DYNAMO_CONFIG
     from torch.export.dynamic_shapes import _process_dynamic_shapes
 
+    log_export_usage(event="export.private_api", flags={"capture_pre_autograd_graph"})
+
     if kwargs is None:
         kwargs = {}
 
diff --git a/torch/export/_trace.py b/torch/export/_trace.py
@@ -286,12 +286,16 @@ def _export_to_torch_ir(
     preserve_module_call_signature: Tuple[str, ...] = (),
     disable_constraint_solver: bool = False,
     restore_fqn: bool = True,
+    _log_export_usage: bool = True,
 ) -> torch.fx.GraphModule:
     """
     Traces either an nn.Module's forward function or just a callable with PyTorch
     operations inside and produce a torch.fx.GraphModule in torch IR.
     """
 
+    if _log_export_usage:
+        log_export_usage(event="export.private_api", flags={"_export_to_torch_ir"})
+
     constraints = constraints or []
     kwargs = kwargs or {}
 
@@ -318,6 +322,7 @@ def _export_to_torch_ir(
                     assume_static_by_default=True,
                     tracing_mode="symbolic",
                     disable_constraint_solver=disable_constraint_solver,
+                    _log_export_usage=_log_export_usage,
                 )(
                     *args,
                     **kwargs,
@@ -580,6 +585,7 @@ def _export(
     log_export_usage(event="export.enter", flags=flags)
 
     if constraints is not None:
+        log_export_usage(event="export.private_api", flags={"constraints"})
         warnings.warn(
             "Using `constraints` to specify dynamic shapes for export is DEPRECATED "
             "and will not be supported in the future. "
@@ -744,6 +750,7 @@ def forward(self, *args, **kwargs):
         constraints,
         preserve_module_call_signature=preserve_module_call_signature,
         restore_fqn=False,  # don't need to restore because we will do it later
+        _log_export_usage=False,
     )
 
     # We detect the fake_mode by looking at gm_torch_level's placeholders, this is the fake_mode created in dynamo.
