add manual NumPy patching to `mkl_fft` by ndgrigorian · Pull Request #224 · IntelPython/mkl_fft

ndgrigorian · 2025-09-10T07:29:27Z

this PR proposes implementing manual patching for NumPy a la mkl_umath to mkl_fft

this patching is specifically designed with NEP-36 in mind

Copilot

Pull request overview

This PR introduces manual NumPy patching functionality to mkl_fft, following NEP-36 fair play guidelines. It allows users to dynamically replace NumPy's FFT functions with mkl_fft's optimized implementations at runtime, similar to the approach used in mkl_umath.

Changes:

Added _patch_numpy.py module with patching/restoration functions and a context manager
Added test file test_patch.py to verify basic patch/restore functionality
Exported new public API functions: patch_numpy_fft(), restore_numpy_fft(), is_patched(), and mkl_fft context manager

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 10 comments.

File	Description
mkl_fft/_patch_numpy.py	Implements the patching mechanism using thread-local storage to track state, provides functions to patch/restore numpy.fft functions, and includes a context manager for temporary patching
mkl_fft/tests/test_patch.py	Basic test verifying patch and restore functionality by checking module names
mkl_fft/init.py	Exports the new patching API functions and context manager class

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

mkl_fft/__init__.py

Copilot · 2026-02-20T06:40:54Z

mkl_fft/_patch_numpy.py

+def patch_numpy_fft(verbose=False):
+    if verbose:
+        print("Now patching NumPy FFT submodule with mkl_fft NumPy interface.")
+        print("Please direct bug reports to https://github.com/IntelPython/mkl_fft")
+    if not _is_tls_initialized():
+        _initialize_tls()
+    _tls.patch.do_patch()


Using thread-local storage for tracking patch state creates a mismatch between local tracking and global effects. When one thread calls patch_numpy_fft(), it modifies the global np.fft namespace (affecting all threads), but only sets _is_patched=True in its own thread-local storage. This means: 1) Thread B won't see that Thread A has patched, and might try to patch again (triggering the double-patch issue), and 2) Thread A might restore while Thread B still needs the patch. Consider using a global lock and counter to properly coordinate multi-threaded patching, or document that this is not thread-safe for concurrent patch/restore calls.

Copilot · 2026-02-20T06:40:54Z

mkl_fft/_patch_numpy.py

+def patch_numpy_fft(verbose=False):
+    if verbose:
+        print("Now patching NumPy FFT submodule with mkl_fft NumPy interface.")
+        print("Please direct bug reports to https://github.com/IntelPython/mkl_fft")
+    if not _is_tls_initialized():
+        _initialize_tls()
+    _tls.patch.do_patch()
+
+
+def restore_numpy_fft(verbose=False):
+    if verbose:
+        print("Now restoring original NumPy FFT submodule.")
+    if not _is_tls_initialized():
+        _initialize_tls()
+    _tls.patch.restore(verbose=verbose)
+
+
+def is_patched():
+    if not _is_tls_initialized():
+        _initialize_tls()
+    return _tls.patch.is_patched()
+
+
+class mkl_fft(ContextDecorator):
+    def __enter__(self):
+        patch_numpy_fft()
+        return self
+
+    def __exit__(self, *exc):
+        restore_numpy_fft()
+        return False


The new patching functionality lacks documentation. Consider adding: 1) docstrings to all public functions (patch_numpy_fft, restore_numpy_fft, is_patched), 2) docstring to the context manager class explaining its usage, 3) updating README.md to explain this new feature and provide usage examples, and 4) adding a note about thread-safety considerations.

Copilot · 2026-02-20T06:40:54Z

mkl_fft/_patch_numpy.py

+    return _tls.patch.is_patched()
+
+
+class mkl_fft(ContextDecorator):


The class name "mkl_fft" conflicts with the module name "mkl_fft". This creates ambiguity: when users import mkl_fft, they get the module, but mkl_fft.init.py also exports a class called mkl_fft. This makes it impossible to use the context manager class directly without confusion. Consider renaming the class to something like "patch_context" or "numpy_fft_context" to avoid this naming collision.

Copilot · 2026-02-20T06:40:54Z

mkl_fft/_patch_numpy.py

+            self._restore_func(name, verbose=verbose)
+        self._is_patched = False
+
+    def do_patch(self):


Calling do_patch() multiple times without calling restore() in between will overwrite _restore_dict values, losing the original NumPy functions. This prevents proper restoration. Add a check: if self._is_patched is True, either skip re-patching or raise an error to prevent this issue.

Suggested change

def do_patch(self):

def do_patch(self):

# Avoid re-patching when already patched, which would overwrite

# the originals stored in _restore_dict with the patched versions.

if self._is_patched:

return

Copilot · 2026-02-20T06:40:55Z

mkl_fft/_patch_numpy.py

+            setattr(np.fft, name, val)
+
+    def restore(self, verbose=False):
+        for name in self._restore_dict.keys():


The restore method iterates over _restore_dict.keys() instead of the original patched_functions list. This means if do_patch() was interrupted partway through (e.g., due to an exception), only the functions that were successfully patched will be restored. While this might be the intended behavior, it could leave numpy in an inconsistent state. Consider whether this is the desired behavior or if all functions should be attempted for restoration.

Suggested change

for name in self._restore_dict.keys():

# Attempt to restore all functions that are part of the patching contract.

# _restore_func will safely handle cases where no original value was recorded.

for name in self.__patched_functions__:

Copilot · 2026-02-20T06:40:55Z

mkl_fft/_patch_numpy.py

+    return (getattr(_tls, "initialized", None) is not None) and (
+        _tls.initialized is True
+    )


The condition check is redundant. The expression "(getattr(_tls, 'initialized', None) is not None) and (_tls.initialized is True)" can be simplified to just "getattr(_tls, 'initialized', False)" since the attribute is always set to True when initialized. The current check is unnecessarily complex.

Suggested change

return (getattr(_tls, "initialized", None) is not None) and (

_tls.initialized is True

)

return getattr(_tls, "initialized", False)

Copilot · 2026-02-20T06:40:55Z

mkl_fft/_patch_numpy.py

+    _is_patched = False
+    __patched_functions__ = _nfft.__all__
+    _restore_dict = {}


Using class attributes instead of instance attributes is a code smell. While each thread gets its own _Patch instance via thread-local storage, using class attributes like _is_patched and _restore_dict suggests they're shared across instances, which is misleading. These should be initialized in an init method as instance attributes (self._is_patched = False, self._restore_dict = {}, self.patched_functions = _nfft.all) to make the code clearer and avoid potential issues if the class is ever used outside of TLS.

Suggested change

_is_patched = False

__patched_functions__ = _nfft.__all__

_restore_dict = {}

def __init__(self):

# Per-instance state for patching; one _Patch instance is stored per-thread.

self._is_patched = False

self.__patched_functions__ = _nfft.__all__

self._restore_dict = {}

Copilot · 2026-02-20T06:40:55Z

mkl_fft/_patch_numpy.py

+
+    def restore(self, verbose=False):
+        for name in self._restore_dict.keys():
+            self._restore_func(name, verbose=verbose)


After restoring all functions, the _restore_dict should be cleared to prevent memory leaks and allow for proper re-patching. Add "self._restore_dict.clear()" after the loop on line 69.

Suggested change

self._restore_func(name, verbose=verbose)

self._restore_func(name, verbose=verbose)

self._restore_dict.clear()

Copilot · 2026-02-20T06:40:56Z

mkl_fft/tests/test_patch.py

+def test_patch():
+    mkl_fft.restore_numpy_fft()
+    assert not mkl_fft.is_patched()
+    assert (np.fft.fft.__module__ == "numpy.fft")
+
+    mkl_fft.patch_numpy_fft()  # Enable mkl_fft in Numpy
+    assert mkl_fft.is_patched()
+    assert (np.fft.fft.__module__ == _nfft.fft.__module__)
+
+    mkl_fft.restore_numpy_fft()  # Disable mkl_fft in Numpy
+    assert not mkl_fft.is_patched()
+    assert (np.fft.fft.__module__ == "numpy.fft")


The test coverage is incomplete. Consider adding tests for: 1) the context manager usage (with mkl_fft.mkl_fft() or similar), 2) testing that patched functions produce correct results, 3) edge cases like double-patching without restoring, 4) thread safety if TLS is being used, and 5) verbose output.

Copilot

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-02-27T10:28:32Z

mkl_fft/_patch_numpy.py

+    def _register_func(self, name, func):
+        if name not in self.__patched_functions__:
+            raise ValueError("%s not an mkl_fft function." % name)
+        f = getattr(np.fft, name)
+        self._restore_dict[name] = f
+        setattr(np.fft, name, func)
+
+    def _restore_func(self, name, verbose=False):
+        if name not in self.__patched_functions__:
+            raise ValueError("%s not an mkl_fft function." % name)
+        try:
+            val = self._restore_dict[name]
+        except KeyError:
+            if verbose:
+                print("failed to restore")
+            return
+        else:
+            if verbose:
+                print("found and restoring...")
+            setattr(np.fft, name, val)
+
+    def restore(self, verbose=False):
+        for name in self._restore_dict.keys():
+            self._restore_func(name, verbose=verbose)
+        self._is_patched = False
+
+    def do_patch(self):
+        for f in self.__patched_functions__:
+            self._register_func(f, getattr(_nfft, f))
+        self._is_patched = True


Patching is not re-entrant/idempotent: calling patch_numpy_fft() twice (or nesting the mkl_fft context manager) overwrites _restore_dict[name] with already-patched functions (getattr(np.fft, name)), so restore_numpy_fft() will not restore NumPy’s original functions. Guard against double-patching (e.g., reference counting / no-op if already patched) and only capture originals once.

Copilot · 2026-02-27T10:28:32Z

mkl_fft/tests/test_patch.py

+    mkl_fft.patch_numpy_fft()  # Enable mkl_fft in Numpy
+    assert mkl_fft.is_patched()
+    assert np.fft.fft.__module__ == _nfft.fft.__module__
+
+    mkl_fft.restore_numpy_fft()  # Disable mkl_fft in Numpy


This test leaves NumPy patched if an assertion fails before restore_numpy_fft() is reached, which can cascade failures into unrelated tests. Wrap the patch/restore portion in try/finally (or use a pytest fixture) to guarantee restoration, and consider adding a regression test for double-patching/nested context usage (currently a broken case).

Suggested change

mkl_fft.patch_numpy_fft() # Enable mkl_fft in Numpy

assert mkl_fft.is_patched()

assert np.fft.fft.__module__ == _nfft.fft.__module__

mkl_fft.restore_numpy_fft() # Disable mkl_fft in Numpy

try:

mkl_fft.patch_numpy_fft() # Enable mkl_fft in Numpy

assert mkl_fft.is_patched()

assert np.fft.fft.__module__ == _nfft.fft.__module__

finally:

mkl_fft.restore_numpy_fft() # Disable mkl_fft in Numpy

Copilot · 2026-02-27T10:28:32Z

mkl_fft/_patch_numpy.py

+class mkl_fft(ContextDecorator):
+    def __enter__(self):
+        patch_numpy_fft()
+        return self
+
+    def __exit__(self, *exc):
+        restore_numpy_fft()
+        return False


Exporting a context manager class named mkl_fft at the package root (from mkl_fft import mkl_fft) is easy to confuse with the mkl_fft module itself. Consider a more explicit name (e.g., numpy_fft_patched/numpy_fft_patch/patched_numpy_fft) to make call sites self-describing.

Copilot · 2026-02-27T10:28:33Z

mkl_fft/_patch_numpy.py

+    _is_patched = False
+    __patched_functions__ = _nfft.__all__
+    _restore_dict = {}
+


_restore_dict is defined as a class attribute (_restore_dict = {}), so it is shared across all _Patch instances (including across threads when used with TLS). This shared mutable state is very likely unintended and can lead to cross-thread interference and hard-to-debug behavior. Make _restore_dict (and _is_patched) instance attributes initialized in __init__ (or otherwise make the singleton nature explicit).

Suggested change

_is_patched = False

__patched_functions__ = _nfft.__all__

_restore_dict = {}

__patched_functions__ = _nfft.__all__

def __init__(self):

self._is_patched = False

self._restore_dict = {}

ndgrigorian changed the title ~~add manual patching to mkl_fft~~ add manual NumPy patching to mkl_fft Sep 10, 2025

ndgrigorian force-pushed the add-mkl-fft-patching branch 3 times, most recently from b2ebfdc to d22ae8e Compare September 10, 2025 16:17

ndgrigorian force-pushed the add-mkl-fft-patching branch from c8fe204 to 559373e Compare September 22, 2025 19:14

ndgrigorian force-pushed the add-mkl-fft-patching branch from 559373e to e4c64d4 Compare October 9, 2025 18:32

ndgrigorian marked this pull request as ready for review February 20, 2026 06:35

ndgrigorian requested a review from antonwolfy as a code owner February 20, 2026 06:35

Copilot AI review requested due to automatic review settings February 20, 2026 06:35

ndgrigorian requested review from jharlow-intel and xaleryb as code owners February 20, 2026 06:35

Copilot started reviewing on behalf of ndgrigorian February 20, 2026 06:36 View session

ndgrigorian added 2 commits February 19, 2026 22:36

add manual patching to mkl_fft

997dfa3

add test for NumPy patching function

b5a4b19

ndgrigorian force-pushed the add-mkl-fft-patching branch from 4eb05ce to b5a4b19 Compare February 20, 2026 06:39

Copilot AI reviewed Feb 20, 2026

View reviewed changes

ndgrigorian force-pushed the add-mkl-fft-patching branch 4 times, most recently from cd55b21 to 51ba5d3 Compare February 20, 2026 09:06

fix patching test

83496b2

ndgrigorian force-pushed the add-mkl-fft-patching branch from 51ba5d3 to 83496b2 Compare February 20, 2026 09:09

ndgrigorian requested a review from Copilot February 27, 2026 10:21

Copilot started reviewing on behalf of ndgrigorian February 27, 2026 10:21 View session

Copilot AI reviewed Feb 27, 2026

View reviewed changes

		return _tls.patch.is_patched()


		class mkl_fft(ContextDecorator):

-    def do_patch(self):
+    def do_patch(self):
+        # Avoid re-patching when already patched, which would overwrite
+        # the originals stored in _restore_dict with the patched versions.
+        if self._is_patched:
+            return

-        for name in self._restore_dict.keys():
+        # Attempt to restore all functions that are part of the patching contract.
+        # _restore_func will safely handle cases where no original value was recorded.
+        for name in self.__patched_functions__:

	self._restore_func(name, verbose=verbose)
	self._restore_func(name, verbose=verbose)
	self._restore_dict.clear()

Conversation

ndgrigorian commented Sep 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Copilot AI Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ndgrigorian commented Sep 10, 2025 •

edited

Loading