Create utility to rebind args/kwargs. (#10987)

Summary:

Refactor to allow users to call `rebind` directly, without the need for running `BindOptionalArgsPass`.

Differential Revision: D75029675

Author: hsharma35

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,31 @@ 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
+
+    sig = matched_schemas[0]
+    bound_args = sig.bind(*args, **kwargs)
+    bound_args.apply_defaults()
+
+    return bound_args.args, bound_args.kwargs