Create utility to rebind args/kwargs.

Differential Revision: D75029675

Pull Request resolved: #10987

Author: hsharma35

.gitignore

@@ -40,3 +40,6 @@ xcuserdata/
 .swiftpm/
 *.xcworkspace/
 *.xcframework/
+
+# Android
+*.aar

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

extension/android/executorch_android/src/main/java/org/pytorch/executorch/Module.java

@@ -47,17 +47,19 @@ public class Module {
    *
    * @param modelPath path to file that contains the serialized ExecuTorch module.
    * @param loadMode load mode for the module. See constants in {@link Module}.
+   * @param numThreads the number of threads to use for inference. A value of 0 defaults to a
+   *     hardware-specific default.
    * @return new {@link org.pytorch.executorch.Module} object which owns the model module.
    */
-  public static Module load(final String modelPath, int loadMode) {
+  public static Module load(final String modelPath, int loadMode, int numThreads = 0) {
     if (!NativeLoader.isInitialized()) {
       NativeLoader.init(new SystemDelegate());
     }
     File modelFile = new File(modelPath);
     if (!modelFile.canRead() || !modelFile.isFile()) {
       throw new RuntimeException("Cannot load model path " + modelPath);
     }
-    return new Module(new NativePeer(modelPath, loadMode));
+    return new Module(new NativePeer(modelPath, loadMode, numThreads));
   }
 
   /**

extension/android/executorch_android/src/main/java/org/pytorch/executorch/NativePeer.java

@@ -28,10 +28,11 @@ class NativePeer {
   private final HybridData mHybridData;
 
   @DoNotStrip
-  private static native HybridData initHybrid(String moduleAbsolutePath, int loadMode);
+  private static native HybridData initHybrid(
+      String moduleAbsolutePath, int loadMode, int numThreads);
 
-  NativePeer(String moduleAbsolutePath, int loadMode) {
-    mHybridData = initHybrid(moduleAbsolutePath, loadMode);
+  NativePeer(String moduleAbsolutePath, int loadMode, int numThreads) {
+    mHybridData = initHybrid(moduleAbsolutePath, loadMode, numThreads);
   }
 
   /** Clean up the native resources associated with this instance */

extension/android/jni/jni_layer.cpp

@@ -228,11 +228,15 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
   static facebook::jni::local_ref<jhybriddata> initHybrid(
       facebook::jni::alias_ref<jclass>,
       facebook::jni::alias_ref<jstring> modelPath,
-      jint loadMode) {
-    return makeCxxInstance(modelPath, loadMode);
+      jint loadMode,
+      jint numThreads) {
+    return makeCxxInstance(modelPath, loadMode, numThreads);
   }
 
-  ExecuTorchJni(facebook::jni::alias_ref<jstring> modelPath, jint loadMode) {
+  ExecuTorchJni(
+      facebook::jni::alias_ref<jstring> modelPath,
+      jint loadMode,
+      jint numThreads) {
     Module::LoadMode load_mode = Module::LoadMode::Mmap;
     if (loadMode == 0) {
       load_mode = Module::LoadMode::File;
@@ -259,11 +263,10 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
     // Based on testing, this is almost universally faster than using all
     // cores, as efficiency cores can be quite slow. In extreme cases, using
     // all cores can be 10x slower than using cores/2.
-    //
-    // TODO Allow overriding this default from Java.
     auto threadpool = executorch::extension::threadpool::get_threadpool();
     if (threadpool) {
-      int thread_count = cpuinfo_get_processors_count() / 2;
+      int thread_count =
+          numThreads != 0 ? numThreads : cpuinfo_get_processors_count() / 2;
       if (thread_count > 0) {
         threadpool->_unsafe_reset_threadpool(thread_count);
       }