Arm backend: Load quantized_ops_aot_lib with arm quantizer module.

Instead of defining a new function for loading the quantized_aot_lib,
and portable_lib, simply import the module to reuse the logic
which is implemented in the __init__ file.

By doing this when importing the arm quantizer, we are sure that
this lib is loaded when running quantized networks.

Signed-off-by: Adrian Lundell <[email protected]>
Change-Id: I5ef585981a09d332d15f7001fb97536b1a56def7

Author: AdrianLundell

backends/arm/quantizer/__init__.py

@@ -13,3 +13,16 @@
 
 # Used in tests
 from .arm_quantizer_utils import is_annotated  # noqa
+
+# Load quantized ops library.
+try:
+    import executorch.extension.pybindings.portable_lib
+    import executorch.kernels.quantized  # noqa
+except:
+    import logging
+
+    logging.info(
+        "Failed to load portable_lib and quantized_aot_lib. To run quantized kernels AOT, either build "
+        "Executorch with pybindings, or load your own custom built op library using torch.ops.load_library."
+    )
+    del logging

backends/arm/test/conftest.py

@@ -3,10 +3,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import glob
 import logging
 import os
-import platform
 import random
 import shutil
 import sys
@@ -81,8 +79,7 @@ def try_addoption(*args, **kwargs):
 
 
 def pytest_sessionstart(session):
-    if not session.config.option.collectonly:
-        _load_libquantized_ops_aot_lib()
+    pass
 
 
 def pytest_sessionfinish(session, exitstatus):
@@ -172,36 +169,3 @@ def get_option(option: str) -> Any | None:
     if option in pytest._test_options:  # type: ignore[attr-defined]
         return pytest._test_options[option]  # type: ignore[attr-defined]
     return None
-
-
-def _load_lib(lib_name_pattern: str):
-    """
-    Find and load a library by name in build_folder.
-    """
-
-    library_paths = glob.glob(lib_name_pattern, recursive=True)
-    if len(library_paths) > 0:
-        import torch
-
-        torch.ops.load_library(library_paths[0])
-    else:
-        raise RuntimeError(
-            f"Did not find any library matching {lib_name_pattern}. Have you installed executorch properly?"
-        )
-
-
-def _load_libquantized_ops_aot_lib():
-    """
-    Find and load the libquantized_ops_aot_lib shared library.
-    """
-
-    so_ext = {
-        "Darwin": "dylib",
-        "Linux": "so",
-        "Windows": "dll",
-    }.get(platform.system(), None)
-
-    executorch_path = os.path.join(os.path.dirname(__file__), "..", "..", "..")
-
-    _load_lib(f"{executorch_path}/**/_portable_lib.cpython-310*{so_ext}")
-    _load_lib(f"{executorch_path}/**/libquantized_ops_aot_lib.{so_ext}")

examples/arm/aot_arm_compiler.py

@@ -29,7 +29,6 @@
     get_symmetric_quantization_config,
     TOSAQuantizer,
 )
-from executorch.backends.arm.test.conftest import _load_libquantized_ops_aot_lib
 from executorch.backends.arm.tosa_partitioner import TOSAPartitioner
 from executorch.backends.arm.tosa_specification import TosaSpecification
 
@@ -510,11 +509,6 @@ def get_args():
     if args.debug:
         logging.basicConfig(level=logging.DEBUG, format=FORMAT, force=True)
 
-    # Load quantized ops library.
-    if args.quantize:
-        logging.info("Loading lib_quantized_custom_op_lib")
-        _load_libquantized_ops_aot_lib()
-
     # if we have custom ops, register them before processing the model
     if args.so_library is not None:
         logging.info(f"Loading custom ops from {args.so_library}")