patching update: is_patched implemented, use_in_numpy, restore, context manager, init

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

@@ -37,12 +37,15 @@ ctypedef struct function_info:
 cdef class patch:
     cdef int functions_count
     cdef function_info* functions
+    cdef bint _is_patched
 
-    functions_dict = {}
+    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:
@@ -95,7 +98,9 @@ cdef class patch:
             index = self.functions_dict[func]
             function = self.functions[index].patch_function
             signature = self.functions[index].signature
-            res = cnp.PyUFunc_ReplaceLoopBySignature(np_umath, function, signature, &temp)
+            res = cnp.PyUFunc_ReplaceLoopBySignature(<cnp.ufunc>np_umath, function, signature, &temp)
+
+        self._is_patched = True
 
     def do_unpatch(self):
         cdef int res
@@ -110,6 +115,10 @@ cdef class patch:
             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()
@@ -123,14 +132,40 @@ def _initialize_tls():
     _tls.patch = patch()
     _tls.initialized = True
 
-    
-def do_patch():
+
+def use_in_numpy():
+    '''
+    Enables using of mkl_umath in Numpy.
+    '''
     if not _is_tls_initialized():
         _initialize_tls()
     _tls.patch.do_patch()
 
 
-def do_unpatch():
+def restore():
+    '''
+    Disables using of mkl_umath in Numpy.
+    '''
     if not _is_tls_initialized():
         _initialize_tls()
     _tls.patch.do_unpatch()
+
+
+def is_patched():
+    '''
+    Returns whether Numpy has been patched with mkl_umath.
+    '''
+    if not _is_tls_initialized():
+        _initialize_tls()
+    _tls.patch.is_patched()
+
+from contextlib import ContextDecorator
+
+class mkl_umath(ContextDecorator):
+    def __enter__(self):
+        use_in_numpy()
+        return self
+
+    def __exit__(self, *exc):
+        restore()
+        return False