Merge pull request #3 from SAT/patching-update

Patching update

Author: Alexander-Makaryev

mkl_umath/__init__.py

@@ -0,0 +1,9 @@
+'''
+Implementation of Numpy universal math functions using Intel(R) MKL and Intel(R) C compiler runtime.
+'''
+
+from ._version import __version__
+
+from ._ufuncs import *
+
+from ._patch import mkl_umath, use_in_numpy, restore, is_patched

mkl_umath/setup.py

@@ -151,8 +151,8 @@ def generate_umath_c(ext, build_dir):
         libraries = mkl_libraries + ['loops_intel'],
         library_dirs = mkl_library_dirs,
         extra_compile_args = [
-            # '-DNDEBUG',
-            '-ggdb', '-O0', '-Wall', '-Wextra', '-DDEBUG',
+            '-DNDEBUG',
+            # '-ggdb', '-O0', '-Wall', '-Wextra', '-DDEBUG',
         ]
     )
 
@@ -168,8 +168,8 @@ def generate_umath_c(ext, build_dir):
         libraries = mkl_libraries + ['loops_intel'],
         library_dirs = mkl_library_dirs,
         extra_compile_args = [
-            # '-DNDEBUG',
-            '-ggdb', '-O0', '-Wall', '-Wextra', '-DDEBUG',
+            '-DNDEBUG',
+            #'-ggdb', '-O0', '-Wall', '-Wextra', '-DDEBUG',
         ]
     )
 

mkl_umath/src/patch.pyx

@@ -25,35 +25,100 @@ import numpy as np
 
 from libc.stdlib cimport malloc, free
 
-cimport cpython.pycapsule
-
 cnp.import_umath()
 
+
 ctypedef struct function_info:
-    cnp.PyUFuncGenericFunction np_function
-    cnp.PyUFuncGenericFunction mkl_function
+    cnp.PyUFuncGenericFunction original_function
+    cnp.PyUFuncGenericFunction patch_function
     int* signature
 
-ctypedef struct functions_struct:
-    int count
-    function_info* functions
-
-
-cdef const char *capsule_name = "functions_cache"
-
-
-cdef void _capsule_destructor(object caps):
-    cdef functions_struct* fs
-
-    if (caps is None):
-        print("Nothing to destroy")
-        return
-    fs = <functions_struct *>cpython.pycapsule.PyCapsule_GetPointer(caps, capsule_name)
-    for i in range(fs[0].count):
-        free(fs[0].functions[i].signature)
-    free(fs[0].functions)
-    free(fs)
 
+cdef class patch:
+    cdef int functions_count
+    cdef function_info* functions
+    cdef bint _is_patched
+
+    functions_dict = dict()
+
+    def __cinit__(self):
+        cdef int pi, oi
+
+        self._is_patched = False
+
+        umaths = [i for i in dir(mu) if isinstance(getattr(mu, i), np.ufunc)]
+        self.functions_count = 0
+        for umath in umaths:
+            mkl_umath = getattr(mu, umath)
+            self.functions_count = self.functions_count + mkl_umath.ntypes
+
+        self.functions = <function_info *> malloc(self.functions_count * sizeof(function_info))
+
+        func_number = 0
+        for umath in umaths:
+            patch_umath = getattr(mu, umath)
+            c_patch_umath = <cnp.ufunc>patch_umath
+            c_orig_umath = <cnp.ufunc>getattr(nu, umath)
+            nargs = c_patch_umath.nargs
+            for pi in range(c_patch_umath.ntypes):
+                oi = 0
+                while oi < c_orig_umath.ntypes:
+                    found = True
+                    for i in range(c_patch_umath.nargs):
+                        if c_patch_umath.types[pi * nargs + i] != c_orig_umath.types[oi * nargs + i]:
+                            found = False
+                            break
+                    if found == True:
+                        break
+                    oi = oi + 1
+                if oi < c_orig_umath.ntypes:
+                    self.functions[func_number].original_function = c_orig_umath.functions[oi]
+                    self.functions[func_number].patch_function = c_patch_umath.functions[pi]
+                    self.functions[func_number].signature = <int *> malloc(nargs * sizeof(int))
+                    for i in range(nargs):
+                        self.functions[func_number].signature[i] = c_patch_umath.types[pi * nargs + i]
+                    self.functions_dict[(umath, patch_umath.types[pi])] = func_number
+                    func_number = func_number + 1
+                else:
+                    raise RuntimeError("Unable to find original function for: " + umath + " " + patch_umath.types[pi])
+
+    def __dealloc__(self):
+        for i in range(self.functions_count):
+            free(self.functions[i].signature)
+        free(self.functions)
+
+    def do_patch(self):
+        cdef int res
+        cdef cnp.PyUFuncGenericFunction temp
+        cdef cnp.PyUFuncGenericFunction function
+        cdef int* signature
+
+        for func in self.functions_dict:
+            np_umath = getattr(nu, func[0])
+            index = self.functions_dict[func]
+            function = self.functions[index].patch_function
+            signature = self.functions[index].signature
+            res = cnp.PyUFunc_ReplaceLoopBySignature(<cnp.ufunc>np_umath, function, signature, &temp)
+
+        self._is_patched = True
+
+    def do_unpatch(self):
+        cdef int res
+        cdef cnp.PyUFuncGenericFunction temp
+        cdef cnp.PyUFuncGenericFunction function
+        cdef int* signature
+
+        for func in self.functions_dict:
+            np_umath = getattr(nu, func[0])
+            index = self.functions_dict[func]
+            function = self.functions[index].original_function
+            signature = self.functions[index].signature
+            res = cnp.PyUFunc_ReplaceLoopBySignature(np_umath, function, signature, &temp)
+
+        self._is_patched = False
+
+    def is_patched(self):
+        return self._is_patched
 
 from threading import local as threading_local
 _tls = threading_local()
@@ -64,103 +129,43 @@ def _is_tls_initialized():
 
 
 def _initialize_tls():
-    cdef functions_struct* fs
-    cdef int funcs_count
-
-    _tls.functions_dict = {}
-
-    umaths = [i for i in dir(mu) if isinstance(getattr(mu, i), np.ufunc)]
-    funcs_count = 0
-    for umath in umaths:
-        mkl_umath = getattr(mu, umath)
-        funcs_count = funcs_count + mkl_umath.ntypes
-
-    fs = <functions_struct *> malloc(sizeof(functions_struct))
-    fs[0].count = funcs_count
-    fs[0].functions = <function_info *> malloc(funcs_count * sizeof(function_info))
-
-    func_number = 0
-    for umath in umaths:
-        mkl_umath = getattr(mu, umath)
-        np_umath = getattr(nu, umath)
-        c_mkl_umath = <cnp.ufunc>mkl_umath
-        c_np_umath = <cnp.ufunc>np_umath
-        for type in mkl_umath.types:
-            np_index = np_umath.types.index(type)
-            fs[0].functions[func_number].np_function = c_np_umath.functions[np_index]
-            mkl_index = mkl_umath.types.index(type)
-            fs[0].functions[func_number].mkl_function = c_mkl_umath.functions[mkl_index]
-
-            nargs = c_mkl_umath.nargs
-            fs[0].functions[func_number].signature = <int *> malloc(nargs * sizeof(int))
-            for i in range(nargs):
-                fs[0].functions[func_number].signature[i] = c_mkl_umath.types[mkl_index*nargs + i]
-
-            _tls.functions_dict[(umath, type)] = func_number
-            func_number = func_number + 1
-
-    _tls.functions_capsule = cpython.pycapsule.PyCapsule_New(<void *>fs, capsule_name, &_capsule_destructor)
-
+    _tls.patch = patch()
     _tls.initialized = True
 
 
-def _get_func_dict():
+def use_in_numpy():
+    '''
+    Enables using of mkl_umath in Numpy.
+    '''
     if not _is_tls_initialized():
         _initialize_tls()
-    return _tls.functions_dict
+    _tls.patch.do_patch()
 
 
-cdef function_info* _get_functions():
-    cdef function_info* functions
-    cdef functions_struct* fs
-
+def restore():
+    '''
+    Disables using of mkl_umath in Numpy.
+    '''
     if not _is_tls_initialized():
         _initialize_tls()
+    _tls.patch.do_unpatch()
 
-    capsule = _tls.functions_capsule
-    if (not cpython.pycapsule.PyCapsule_IsValid(capsule, capsule_name)):
-        raise ValueError("Internal Error: invalid capsule stored in TLS")
-    fs = <functions_struct *>cpython.pycapsule.PyCapsule_GetPointer(capsule, capsule_name)
-    return fs[0].functions
-
-
-cdef void c_do_patch():
-    cdef int res
-    cdef cnp.PyUFuncGenericFunction temp
-    cdef cnp.PyUFuncGenericFunction function
-    cdef int* signature
-
-    funcs_dict = _get_func_dict()
-    functions = _get_functions()
-
-    for func in funcs_dict:
-        np_umath = getattr(nu, func[0])
-        index = funcs_dict[func]
-        function = functions[index].mkl_function
-        signature = functions[index].signature
-        res = cnp.PyUFunc_ReplaceLoopBySignature(np_umath, function, signature, &temp)
-
-
-cdef void c_do_unpatch():
-    cdef int res
-    cdef cnp.PyUFuncGenericFunction temp
-    cdef cnp.PyUFuncGenericFunction function
-    cdef int* signature
-
-    funcs_dict = _get_func_dict()
-    functions = _get_functions()
-
-    for func in funcs_dict:
-        np_umath = getattr(nu, func[0])
-        index = funcs_dict[func]
-        function = functions[index].np_function
-        signature = functions[index].signature
-        res = cnp.PyUFunc_ReplaceLoopBySignature(np_umath, function, signature, &temp)
 
+def is_patched():
+    '''
+    Returns whether Numpy has been patched with mkl_umath.
+    '''
+    if not _is_tls_initialized():
+        _initialize_tls()
+    _tls.patch.is_patched()
 
-def do_patch():
-    c_do_patch()
+from contextlib import ContextDecorator
 
+class mkl_umath(ContextDecorator):
+    def __enter__(self):
+        use_in_numpy()
+        return self
 
-def do_unpatch():
-    c_do_unpatch()
+    def __exit__(self, *exc):
+        restore()
+        return False