feat(framework): add cython wrappers mode (#27770)

Co-authored-by: ivy-branch <[email protected]>

Author: yopknopixx

.gitignore

@@ -21,6 +21,7 @@ with_time_logs/
 *.jpg
 *.jpeg
 *.gif
+*.so
 .hypothesis
 .array_api_tests_k_flag*
 internal_automation_tools/

ivy/__init__.py

@@ -591,7 +591,7 @@ class Node(str):
 nan_policy_stack = []
 dynamic_backend_stack = []
 warn_to_regex = {"all": "!.*", "ivy_only": "^(?!.*ivy).*$", "none": ".*"}
-
+cython_wrappers_stack = []
 
 # local
 import threading
@@ -741,7 +741,7 @@ class Node(str):
 
 locks = {"backend_setter": threading.Lock()}
 
-
+from .wrappers import *
 from .func_wrapper import *
 from .data_classes.array import Array, add_ivy_array_instance_methods
 from .data_classes.array.conversions import *
@@ -776,6 +776,7 @@ class Node(str):
     choose_random_backend,
     unset_backend,
 )
+from . import wrappers
 from . import func_wrapper
 from .utils import assertions, exceptions, verbosity
 from .utils.backend import handler
@@ -961,6 +962,7 @@ def __deepcopy__(self, memo):
     "default_uint_dtype_stack": data_type.default_uint_dtype_stack,
     "nan_policy_stack": nan_policy_stack,
     "dynamic_backend_stack": dynamic_backend_stack,
+    "cython_wrappers_stack": cython_wrappers_stack,
 })
 
 _default_globals = copy.deepcopy(globals_vars)
@@ -1144,7 +1146,7 @@ def unset_nan_policy():
 ivy.dynamic_backend = dynamic_backend_stack[-1] if dynamic_backend_stack else True
 
 
-def set_dynamic_backend(flag):
+def set_dynamic_backend(flag):  # noqa: D209
     """Set the global dynamic backend setting to the provided flag (True or
     False)"""
     global dynamic_backend_stack
@@ -1166,6 +1168,37 @@ def unset_dynamic_backend():
         ivy.__setattr__("dynamic_backend", flag, True)
 
 
+# Cython wrappers
+
+ivy.cython_wrappers_mode = cython_wrappers_stack[-1] if cython_wrappers_stack else False
+
+
+@handle_exceptions
+def set_cython_wrappers_mode(flag: bool = True) -> None:
+    """Set the mode of whether to use cython wrappers for functions.
+
+    Parameter
+    ---------
+    flag
+        boolean whether to use cython wrappers for functions
+
+    Examples
+    --------
+    >>> ivy.set_cython_wrappers_mode(False)
+    >>> ivy.cython_wrappers_mode
+    False
+
+    >>> ivy.set_cython_wrappers_mode(True)
+    >>> ivy.cython_wrappers_mode
+    True
+    """
+    global cython_wrappers_stack
+    if flag not in [True, False]:
+        raise ValueError("cython_wrappers_mode must be a boolean value (True or False)")
+    cython_wrappers_stack.append(flag)
+    ivy.__setattr__("cython_wrappers_mode", flag, True)
+
+
 # Context Managers
 
 
@@ -1438,6 +1471,7 @@ def cast_data_types(val=True):
     "default_int_dtype",
     "default_complex_dtype",
     "default_uint_dtype",
+    "cython_wrappers_mode",
 ]
 
 
@@ -1479,7 +1513,7 @@ def set_logging_mode(self, mode):
         logging.getLogger().setLevel(mode)
         self.logging_mode_stack.append(mode)
 
-    def unset_logging_mode(self):
+    def unset_logging_mode(self):  # noqa: D209
         """Remove the most recently set logging mode, returning to the previous
         one."""
         if len(self.logging_mode_stack) > 1:

ivy/func_wrapper.py

@@ -1002,6 +1002,29 @@ def _to_ivy_array(x):
     return _temp_asarray_wrapper
 
 
+# Download compiled cython wrapper wrapper
+
+
+def download_cython_wrapper_wrapper(fn: Callable) -> Callable:
+    @functools.wraps(fn)
+    def _download_cython_wrapper_wrapper(*args, **kwargs):
+        """Wrap the function to download compiled cython wrapper for the
+        function and re-wraps it with the downloaded wrapper.
+
+        Download the compiled cython wrapper by calling
+        ivy.wrappers.get_wrapper(func_name: str) and then wrap the
+        function with the downloaded wrapper.
+        """
+        ivy.wrappers.download_cython_wrapper(fn.__name__)
+        ivy.wrappers.load_one_wrapper(fn.__name__)
+        ivy.functional.__dict__[fn.__name__] = getattr(
+            ivy.wrappers, fn.__name__ + "_wrapper"
+        )(fn)
+        return ivy.functional.__dict__[fn.__name__](*args, **kwargs)
+
+    return _download_cython_wrapper_wrapper
+
+
 # Functions #
 
 
@@ -1046,6 +1069,12 @@ def _wrap_function(
                 )
         return to_wrap
     if isinstance(to_wrap, FunctionType):
+        if ivy.cython_wrappers_mode and ivy.wrappers.wrapper_exists(to_wrap.__name__):
+            if to_wrap.__name__ + "_wrapper" in ivy.wrappers.__all__:
+                to_wrap = getattr(ivy.wrappers, to_wrap.__name__ + "_wrapper")(to_wrap)
+                return to_wrap
+            else:
+                return download_cython_wrapper_wrapper(to_wrap)
         # set attributes
         for attr in original.__dict__.keys():
             # private attribute or decorator

ivy/utils/_importlib.py

@@ -12,7 +12,7 @@
 # expected. Import these modules along with Ivy initialization, as the import logic
 # assumes they exist in sys.modules.
 
-MODULES_TO_SKIP = ["ivy.compiler", "ivy.engines"]
+MODULES_TO_SKIP = ["ivy.compiler", "ivy.engines", "ivy.wrappers"]
 
 IS_COMPILING_WITH_BACKEND = False
 

ivy/utils/backend/handler.py

@@ -358,7 +358,6 @@ def set_backend(backend: str, dynamic: bool = False):
         global ivy_original_dict
         if not backend_stack:
             ivy_original_dict = ivy.__dict__.copy()
-
         _clear_current_sub_backends()
         if isinstance(backend, str):
             temp_stack = []

ivy/wrappers/__init__.py

@@ -0,0 +1,29 @@
+import os
+import sys
+import glob
+import importlib
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+so_files = glob.glob(dir_path + "/*.so")
+sys.path.append(dir_path)
+
+__all__ = []
+
+for so_file in so_files:
+    # if os.path.basename(so_file) != "add.so":
+    #     continue
+    module_name = os.path.splitext(os.path.basename(so_file))[0]
+
+    locals()[module_name] = importlib.import_module(module_name)
+
+    if module_name + "_wrapper" in locals()[module_name].__dict__.keys():
+        locals()[module_name + "_wrapper"] = getattr(
+            locals()[module_name], module_name + "_wrapper"
+        )
+        __all__.append(module_name + "_wrapper")
+
+del dir_path
+del so_files
+
+import utils
+from utils import *

ivy/wrappers/utils.py

@@ -0,0 +1,54 @@
+import os
+import logging
+import json
+from urllib import request
+import importlib
+import ivy
+
+folder_path = os.sep.join(__file__.split(os.sep)[:-3])
+wrappers_path = os.path.join(folder_path, "wrappers.json")
+wrappers = json.loads(open(wrappers_path).read())
+wrapers_dir = os.path.join(folder_path, "ivy/wrappers")
+
+
+def download_cython_wrapper(func_name: str):
+    """Get the wrapper for the given function name."""
+    if func_name + ".so" not in wrappers["ivy"]["functional"]:
+        logging.warn(f"Wrapper for {func_name} not found.")
+        return False
+    try:
+        response = request.urlopen(
+            "https://raw.githubusercontent.com/unifyai"
+            + "/binaries/cython_wrappers/wrappers/"
+            + func_name
+            + ".so"
+        )
+        os.makedirs(wrapers_dir, exist_ok=True)
+        with open(os.path.join(wrapers_dir, func_name + ".so"), "wb") as f:
+            f.write(response.read())
+        print("Downloaded wrapper for " + func_name)
+        return True
+    except request.HTTPError:
+        logging.warn(f"Unable to download wrapper for {func_name}.")
+        return False
+
+
+def wrapper_exists(func_name: str):
+    """Check if the wrapper for the given function name exists."""
+    return func_name + ".so" in wrappers["ivy"]["functional"]
+
+
+def load_one_wrapper(func_name: str):
+    """Load the wrapper for the given function name."""
+    module_name = func_name
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+    # check if file exists
+    if os.path.isfile(os.path.join(dir_path, module_name + ".so")):
+        ivy.wrappers.__dict__[module_name] = importlib.import_module(module_name)
+        ivy.wrappers.__dict__[module_name + "_wrapper"] = getattr(
+            ivy.wrappers.__dict__[module_name], module_name + "_wrapper"
+        )
+        ivy.wrappers.__all__.append(module_name + "_wrapper")
+        return True
+    else:
+        return False

ivy_tests/test_ivy/conftest.py

@@ -292,6 +292,10 @@ def process_cl_flags(config) -> Dict[str, bool]:
             False,
             getopt("--with-transpile"),
         ),
+        "test_cython_wrapper": (
+            getopt("--skip-cython-wrapper-testing"),
+            getopt("--with-cython-wrapper-testing"),
+        ),
     }
 
     # whether to skip gt testing or not
@@ -358,6 +362,8 @@ def pytest_addoption(parser):
         default=None,
         help="Print test items in my custom format",
     )
+    parser.addoption("--skip-cython-wrapper-testing", action="store_true")
+    parser.addoption("--with-cython-wrapper-testing", action="store_true")
 
 
 def pytest_collection_finish(session):

ivy_tests/test_ivy/helpers/function_testing.py

@@ -154,6 +154,11 @@ def test_function_backend_computation(
             test_flags.container[0] for _ in range(total_num_arrays)
         ]
 
+    if test_flags.test_cython_wrapper:
+        ivy.set_cython_wrappers_mode(True)
+    else:
+        ivy.set_cython_wrappers_mode(False)
+
     with BackendHandler.update_backend(fw) as ivy_backend:
         # Update variable flags to be compatible with float dtype and with_out args
         test_flags.as_variable = [

ivy_tests/test_ivy/helpers/pipeline_helper.py

@@ -11,11 +11,12 @@ class BackendHandlerMode(Enum):
 
 
 class WithBackendContext:
-    def __init__(self, backend) -> None:
+    def __init__(self, backend, cached=True) -> None:
         self.backend = backend
+        self.cached = cached
 
     def __enter__(self):
-        return ivy.with_backend(self.backend)
+        return ivy.with_backend(self.backend, cached=self.cached)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         return