Merge branch 'main' into migrate-pt2e-top-level

Author: metascroy

CMakePresets.json

@@ -4,8 +4,7 @@
     {
       "name": "common",
       "hidden": true,
-      "binaryDir": "${sourceDir}/cmake-out",
-      "generator": "Unix Makefiles"
+      "binaryDir": "${sourceDir}/cmake-out"
     },
     {
       "name": "macos-arm64",

backends/cadence/aot/TARGETS

@@ -211,6 +211,7 @@ python_library(
     typing = True,
     deps = [
         ":pass_utils",
+        ":utils",
         "//executorch/backends/cadence/aot:pass_utils",
         "//executorch/exir:pass_base",
         "//executorch/exir/dialects:lib",

backends/cadence/aot/simplify_ops.py

@@ -16,10 +16,10 @@
     CadencePassAttribute,
     register_cadence_pass,
 )
+from executorch.backends.cadence.aot.utils import rebind
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.pass_base import ExportPass, ProxyValue
-from torch.fx.operator_schemas import get_signature_for_torch_op
 
 
 @register_cadence_pass(CadencePassAttribute(opt_level=0))
@@ -117,32 +117,11 @@ class BindOptionalArgsPass(ExportPass):
     def call_operator(self, op, args, kwargs, meta):
         if not isinstance(op, EdgeOpOverload):
             return super().call_operator(op, args, kwargs, meta)
-        assert callable(op)
 
-        torch_op_schemas = get_signature_for_torch_op(op._op)
-        if len(torch_op_schemas) == 0:
-            return super().call_operator(op, args, kwargs, meta)
-
-        matched_schemas = []
-        # Iterate through all of the schema until we find one that matches
-        # If one matches, populate `new_args_and_kwargs` with the new args/kwargs
-        # values. If none matches, `new_args_and_kwargs` will be None
-        for candidate_signature in torch_op_schemas:
-            try:
-                candidate_signature.bind(*args, **kwargs)
-                matched_schemas.append(candidate_signature)
-            except TypeError:
-                continue
-
-        if len(matched_schemas) != 1:
-            # Did not match any schema. Cannot normalize
-            return super().call_operator(op, args, kwargs, meta)
-
-        sig = matched_schemas[0]
-        bound_args = sig.bind(*args, **kwargs)
-        bound_args.apply_defaults()
+        if (updated_args := rebind(op, args, kwargs)) is not None:
+            args, kwargs = updated_args
 
-        return super().call_operator(op, bound_args.args, bound_args.kwargs, meta)
+        return super().call_operator(op, args, kwargs, meta)
 
 
 # This class encapsulates all the functions that simplify the op's args

backends/cadence/aot/utils.py

@@ -18,7 +18,9 @@
 from executorch.exir import ExecutorchProgramManager, memory
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload, EdgeOpOverloadPacket
+from executorch.exir.pass_base import Argument
 from tabulate import tabulate
+from torch.fx.operator_schemas import get_signature_for_torch_op
 
 from torch.utils._pytree import tree_flatten
 
@@ -308,3 +310,30 @@ def get_size(self, exir_id: int) -> int:
 # Return default memory config for the backend
 def get_default_memory_config() -> MemoryConfig:
     return MemoryConfig(memory_sizes=[0x1000000000])
+
+
+def rebind(
+    op: EdgeOpOverload, args: tuple[Argument, ...], kwargs: dict[str, Argument]
+) -> Optional[tuple[tuple[Argument, ...], dict[str, Argument]]]:
+    """Populates optional args and binds args/kwargs based on schema."""
+    torch_op_schemas = get_signature_for_torch_op(op._op)
+
+    matched_schemas = []
+    # Iterate through all of the schema until we find one that matches
+    # If one matches, populate `new_args_and_kwargs` with the new args/kwargs
+    # values. If none matches, `new_args_and_kwargs` will be None
+    for candidate_signature in torch_op_schemas:
+        try:
+            candidate_signature.bind(*args, **kwargs)
+            matched_schemas.append(candidate_signature)
+        except TypeError:
+            continue
+
+    if len(matched_schemas) != 1:
+        # Did not match any schema. Cannot normalize
+        return None
+
+    bound_args = matched_schemas[0].bind(*args, **kwargs)
+    bound_args.apply_defaults()
+
+    return bound_args.args, bound_args.kwargs

devtools/inspector/_inspector.py

@@ -1240,7 +1240,7 @@ def find_total_for_module(self, module_name: str) -> float:
         for block in self.event_blocks:
             for event in block.events:
                 # Skip OPERATOR_CALL events to avoid double-counting and exclude framework tax
-                if event.event_name == "OPERATOR_CALL":
+                if event.name == "OPERATOR_CALL":
                     continue
 
                 module_hierarchy = event.module_hierarchy.values()

docs/source/_static/img/inspector_tabular_output.png

@@ -187,24 +187,22 @@ def forward(self, x):
 
 from executorch.devtools import Inspector
 
+
 # sphinx_gallery_start_ignore
-inspector_patch = patch.object(Inspector, "__init__", return_value=None)
-inspector_patch_print = patch.object(Inspector, "print_data_tabular", return_value="")
-inspector_patch.start()
+# inspector_patch = patch.object(Inspector, "__init__", return_value=None)
+inspector_patch_print = patch.object(Inspector, "print_data_tabular", return_value=None)
 inspector_patch_print.start()
 # sphinx_gallery_end_ignore
 etrecord_path = "etrecord.bin"
 etdump_path = "etdump.etdp"
 inspector = Inspector(etdump_path=etdump_path, etrecord=etrecord_path)
-# sphinx_gallery_start_ignore
-inspector.event_blocks = []
-# sphinx_gallery_end_ignore
 inspector.print_data_tabular()
 
-# sphinx_gallery_start_ignore
-inspector_patch.stop()
-inspector_patch_print.stop()
-# sphinx_gallery_end_ignore
+####################################
+#
+# Here is an example output:
+#
+# .. image:: ../_static/img/inspector_tabular_output.png
 
 ######################################################################
 # Analyzing with an Inspector
@@ -234,11 +232,13 @@ def forward(self, x):
         if event.name == "native_call_addmm.out":
             print(event.name, event.perf_data.raw if event.perf_data else "")
 
+    print()
     # Via Dataframe
     df = event_block.to_dataframe()
     df = df[df.event_name == "native_call_addmm.out"]
-    print(df[["event_name", "raw"]])
-    print()
+    if len(df) > 0:
+        print(df[["event_name", "raw"]])
+        print()
 
 ######################################################################
 # If a user wants to trace an operator back to their model code, they would do
@@ -255,6 +255,20 @@ def forward(self, x):
     if slowest is not None:
         print(slowest.name)
         print()
+        # sphinx_gallery_start_ignore
+        slowest_print = patch.object(
+            slowest,
+            "stack_traces",
+            new={
+                "aten_convolution_default_1_": "  File "
+                '"devtools-integration-tutorial.py", '
+                "line 82, in forward\n"
+                "    x = F.max_pool2d(F.relu(self.conv2(x)), "
+                "2)\n"
+            },
+        )
+        slowest_print.start()
+        # sphinx_gallery_end_ignore
         pp.pprint(slowest.stack_traces)
         print()
         pp.pprint(slowest.module_hierarchy)
@@ -264,7 +278,7 @@ def forward(self, x):
     df = df[df.event_name == "native_call_convolution.out"]
     if len(df) > 0:
         slowest = df.loc[df["p50"].idxmax()]
-        assert slowest
+        assert slowest is not None
         print(slowest.name)
         print()
         pp.pprint(slowest.stack_traces if slowest.stack_traces else "")
@@ -281,6 +295,7 @@ def forward(self, x):
 ######################################################################
 # Note: ``find_total_for_module`` is a special first class method of
 # `Inspector <../model-inspector.html>`__
+#
 
 ######################################################################
 # Conclusion