Move dpnp monkey patching to patch file

Author: ZzEeKkAa

numba_dpex/__init__.py

@@ -17,7 +17,7 @@
 from numba import __version__ as numba_version
 
 from .kernel_api_impl.spirv import target as spirv_kernel_target
-from .numba_patches import patch_is_ufunc
+from .numba_patches import patch_ufuncs
 from .register_kernel_api_overloads import init_kernel_api_spirv_overloads
 
 
@@ -70,7 +70,7 @@ def parse_sem_version(version_string: str) -> Tuple[int, int, int]:
 dpctl_sem_version = parse_sem_version(dpctl.__version__)
 
 # Monkey patches
-patch_is_ufunc.patch()
+patch_ufuncs.patch()
 
 from numba import prange  # noqa E402
 

numba_dpex/core/typing/dpnpdecl.py

@@ -2,10 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import logging
-
 import dpnp
-import numpy as np
 from numba.core import types
 from numba.core.typing.npydecl import (
     Numpy_rules_ufunc,
@@ -25,7 +22,7 @@ def _install_operations(cls: NumpyRulesArrayOperator):
     for op, ufunc_name in cls._op_map.items():
         infer_global(op)(
             type(
-                "NumpyRulesArrayOperator_" + ufunc_name,
+                "DpnpRulesArrayOperator_" + ufunc_name,
                 (cls,),
                 dict(key=op),
             )
@@ -35,36 +32,7 @@ def _install_operations(cls: NumpyRulesArrayOperator):
 class DpnpRulesArrayOperator(NumpyRulesArrayOperator):
     @property
     def ufunc(self):
-        try:
-            dpnpop = getattr(dpnp, self._op_map[self.key])
-            npop = getattr(np, self._op_map[self.key])
-            if not hasattr(dpnpop, "nin"):
-                dpnpop.nin = npop.nin
-            if not hasattr(dpnpop, "nout"):
-                dpnpop.nout = npop.nout
-            if not hasattr(dpnpop, "nargs"):
-                dpnpop.nargs = dpnpop.nin + dpnpop.nout
-
-            # Check if the dpnp operation has a `types` attribute and if an
-            # AttributeError gets raised then "monkey patch" the attribute from
-            # numpy. If the attribute lookup raised a ValueError, it indicates
-            # that dpnp could not be resolve the supported types for the
-            # operation. Dpnp will fail to resolve the `types` if no SYCL
-            # devices are available on the system. For such a scenario, we print
-            # a user-level warning.
-            try:
-                dpnpop.types
-            except ValueError:
-                logging.exception(
-                    f"The types attribute for the {dpnpop} fuction could not "
-                    "be determined."
-                )
-            except AttributeError:
-                dpnpop.types = npop.types
-            dpnpop.is_dpnp_ufunc = True
-            return dpnpop
-        except:
-            pass
+        return getattr(dpnp, self._op_map[self.key])
 
     @classmethod
     def install_operations(cls):

numba_dpex/dpnp_iface/dpnp_ufunc_db.py

@@ -3,7 +3,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import copy
-import logging
 
 import dpnp
 import numpy as np
@@ -57,15 +56,9 @@ def _fill_ufunc_db_with_dpnp_ufuncs(ufunc_db):
     # variable is passed by value
     from numba.np.ufunc_db import _ufunc_db
 
-    failed_dpnpop_types_lst = []
     for ufuncop in dpnpdecl.supported_ufuncs:
         if ufuncop == "erf":
             op = getattr(dpnp, "erf")
-            op.nin = 1
-            op.nout = 1
-            op.nargs = 2
-            op.types = ["f->f", "d->d"]
-            op.is_dpnp_ufunc = True
 
             _unary_d_d = types.float64(types.float64)
             _unary_f_f = types.float32(types.float32)
@@ -76,31 +69,7 @@ def _fill_ufunc_db_with_dpnp_ufuncs(ufunc_db):
         else:
             dpnpop = getattr(dpnp, ufuncop)
             npop = getattr(np, ufuncop)
-            if not hasattr(dpnpop, "nin"):
-                dpnpop.nin = npop.nin
-            if not hasattr(dpnpop, "nout"):
-                dpnpop.nout = npop.nout
-            if not hasattr(dpnpop, "nargs"):
-                dpnpop.nargs = dpnpop.nin + dpnpop.nout
-
-            # Check if the dpnp operation has a `types` attribute and if an
-            # AttributeError gets raised then "monkey patch" the attribute from
-            # numpy. If the attribute lookup raised a ValueError, it indicates
-            # that dpnp could not be resolve the supported types for the
-            # operation. Dpnp will fail to resolve the `types` if no SYCL
-            # devices are available on the system. For such a scenario, we log
-            # dpnp operations for which the ValueError happened and print them
-            # as a user-level warning. It is done this way so that the failure
-            # to load the dpnpdecl registry due to the ValueError does not
-            # impede a user from importing numba-dpex.
-            try:
-                dpnpop.types
-            except ValueError:
-                failed_dpnpop_types_lst.append(ufuncop)
-            except AttributeError:
-                dpnpop.types = npop.types
-
-            dpnpop.is_dpnp_ufunc = True
+
             cp = copy.copy(_ufunc_db[npop])
             ufunc_db.update({dpnpop: cp})
             for key in list(ufunc_db[dpnpop].keys()):
@@ -111,13 +80,3 @@ def _fill_ufunc_db_with_dpnp_ufuncs(ufunc_db):
                     or "D->" in key
                 ):
                     ufunc_db[dpnpop].pop(key)
-
-    if failed_dpnpop_types_lst:
-        try:
-            getattr(dpnp, failed_dpnpop_types_lst[0]).types
-        except ValueError:
-            ops = " ".join(failed_dpnpop_types_lst)
-            logging.exception(
-                "The types attribute for the following dpnp ops could not be "
-                f"determined: {ops}"
-            )

numba_dpex/numba_patches/patch_is_ufunc.py

@@ -0,0 +1,94 @@
+# SPDX-FileCopyrightText: 2020 - 2024 Intel Corporation
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import dpnp
+import numpy as np
+
+from numba_dpex.core.typing import dpnpdecl
+
+
+def patch():
+    patch_is_ufunc()
+    patch_ufuncs()
+
+
+def patch_is_ufunc():
+    """Patches the numba.np.ufunc.array_exprs._is_ufunc function to make it
+    possible to support dpnp universal functions (ufuncs).
+
+    The extra condition is the check for the "is_dpnp_ufunc" attribute to
+    identify a non-NumPy ufunc.
+    """
+    import numpy
+    from numba.np.ufunc.dufunc import DUFunc
+
+    def _is_ufunc(func):
+        return isinstance(func, (numpy.ufunc, DUFunc)) or hasattr(
+            func, "is_dpnp_ufunc"
+        )
+
+    from numba.np.ufunc import array_exprs
+
+    array_exprs._is_ufunc = _is_ufunc
+
+
+def patch_ufuncs():
+    """Patches dpnp user functions to make them compatible with numpy, so we
+    can reuse numba's implementation.
+
+    It adds "nin", "nout", "nargs" and "is_dpnp_ufunc" attributes to ufuncs.
+    """
+    failed_dpnpop_types_lst = []
+
+    op = getattr(dpnp, "erf")
+    op.nin = 1
+    op.nout = 1
+    op.nargs = 2
+    op.types = ["f->f", "d->d"]
+    op.is_dpnp_ufunc = True
+
+    for ufuncop in dpnpdecl.supported_ufuncs:
+        if ufuncop == "erf":
+            continue
+
+        dpnpop = getattr(dpnp, ufuncop)
+        npop = getattr(np, ufuncop)
+
+        if not hasattr(dpnpop, "nin"):
+            dpnpop.nin = npop.nin
+        if not hasattr(dpnpop, "nout"):
+            dpnpop.nout = npop.nout
+        if not hasattr(dpnpop, "nargs"):
+            dpnpop.nargs = dpnpop.nin + dpnpop.nout
+
+        # Check if the dpnp operation has a `types` attribute and if an
+        # AttributeError gets raised then "monkey patch" the attribute from
+        # numpy. If the attribute lookup raised a ValueError, it indicates
+        # that dpnp could not be resolve the supported types for the
+        # operation. Dpnp will fail to resolve the `types` if no SYCL
+        # devices are available on the system. For such a scenario, we log
+        # dpnp operations for which the ValueError happened and print them
+        # as a user-level warning. It is done this way so that the failure
+        # to load the dpnpdecl registry due to the ValueError does not
+        # impede a user from importing numba-dpex.
+        try:
+            dpnpop.types
+        except ValueError:
+            failed_dpnpop_types_lst.append(ufuncop)
+        except AttributeError:
+            dpnpop.types = npop.types
+
+        dpnpop.is_dpnp_ufunc = True
+
+    if len(failed_dpnpop_types_lst) > 0:
+        try:
+            getattr(dpnp, failed_dpnpop_types_lst[0]).types
+        except ValueError:
+            ops = " ".join(failed_dpnpop_types_lst)
+            logging.exception(
+                "The types attribute for the following dpnp ops could not be "
+                f"determined: {ops}"
+            )