Fix GIL for fft extensions

CHANGELOG.md

@@ -86,6 +86,7 @@ This release is compatible with NumPy 2.5.
 * Released the GIL before the remaining blocking OneMKL BLAS and LAPACK calls to prevent host tasks contention, completing the work started in [#2850](https://github.com/IntelPython/dpnp/pull/2850) [#3027](https://github.com/IntelPython/dpnp/pull/3027)
 * Fixed `dpnp.repeat` raising an unclear `TypeError` for a nested sequence of `repeats` [#3024](https://github.com/IntelPython/dpnp/pull/3024)
 * Fixed `dpnp.ndarray.view` ignoring the USM element offset of a sliced array, which also caused `dpnp.einsum` to silently return wrong results for a single sliced operand with no summed index [#3037](https://github.com/IntelPython/dpnp/pull/3037)
+* Released the GIL before the blocking OneMKL DFT calls in the FFT extension [#3040](https://github.com/IntelPython/dpnp/pull/3040)
 
 ### Security
 

dpnp/backend/extensions/fft/common.hpp

@@ -61,7 +61,13 @@ class DescriptorWrapper
                                   "device does not support double precision.");
         }
 
-        descr_.commit(q);
+        {
+            // Release GIL to avoid serialization of host task submissions
+            // to the same queue in OneMKL
+            py::gil_scoped_release lock{};
+
+            descr_.commit(q);
+        }
         queue_ptr_ = std::make_unique<sycl::queue>(q);
     }
 

dpnp/backend/extensions/fft/in_place.tpp

@@ -83,13 +83,19 @@ std::pair<sycl::event, sycl::event>
     // in-place is only used for c2c FFT at this time, passing true or false is
     // indifferent
     using ScaleT = typename ScaleType<prec, dom, true>::type_in;
+    // get_data() calls into the Python C-API and so must be called while the
+    // GIL is still held
     ScaleT *in_out_ptr = in_out.get_data<ScaleT>();
 
     sycl::event fft_event = {};
     std::stringstream error_msg;
     bool is_exception_caught = false;
 
     try {
+        // Release GIL to avoid serialization of host task submissions
+        // to the same queue in OneMKL
+        py::gil_scoped_release lock{};
+
         if (is_forward) {
             fft_event = mkl_dft::compute_forward(descr.get_descriptor(),
                                                  in_out_ptr, depends);

dpnp/backend/extensions/fft/out_of_place.tpp

@@ -138,24 +138,34 @@ std::pair<sycl::event, sycl::event>
     dpnp::tensor::validation::CheckWritable::throw_if_not_writable(out);
     dpnp::tensor::validation::AmpleMemory::throw_if_not_ample(out, n_elems);
 
+    // the input and output types depend on is_forward, so the untyped
+    // pointers are cast to the expected type below. get_data() calls into the
+    // Python C-API and so must be called while the GIL is still held
+    char *in_data = in.get_data();
+    char *out_data = out.get_data();
+
     sycl::event fft_event = {};
     std::stringstream error_msg;
     bool is_exception_caught = false;
 
     try {
+        // Release GIL to avoid serialization of host task submissions
+        // to the same queue in OneMKL
+        py::gil_scoped_release lock{};
+
         if (is_forward) {
             using ScaleT_in = typename ScaleType<prec, dom, true>::type_in;
             using ScaleT_out = typename ScaleType<prec, dom, true>::type_out;
-            ScaleT_in *in_ptr = in.get_data<ScaleT_in>();
-            ScaleT_out *out_ptr = out.get_data<ScaleT_out>();
+            ScaleT_in *in_ptr = reinterpret_cast<ScaleT_in *>(in_data);
+            ScaleT_out *out_ptr = reinterpret_cast<ScaleT_out *>(out_data);
             fft_event = mkl_dft::compute_forward(descr.get_descriptor(), in_ptr,
                                                  out_ptr, depends);
         }
         else {
             using ScaleT_in = typename ScaleType<prec, dom, false>::type_in;
             using ScaleT_out = typename ScaleType<prec, dom, false>::type_out;
-            ScaleT_in *in_ptr = in.get_data<ScaleT_in>();
-            ScaleT_out *out_ptr = out.get_data<ScaleT_out>();
+            ScaleT_in *in_ptr = reinterpret_cast<ScaleT_in *>(in_data);
+            ScaleT_out *out_ptr = reinterpret_cast<ScaleT_out *>(out_data);
             fft_event = mkl_dft::compute_backward(descr.get_descriptor(),
                                                   in_ptr, out_ptr, depends);
         }

dpnp/tests/test_fft_gil.py

@@ -0,0 +1,115 @@
+"""Blocking oneMKL calls in the FFT extension must release the GIL.
+
+Progress of a competing Python thread is measured while the FFT call blocks,
+and compared against ``time.sleep()`` of the same duration, which is ``nogil``
+and therefore the best rate achievable on the machine.
+"""
+
+import sys
+import threading
+import time
+
+import pytest
+
+import dpnp
+
+from .helper import has_support_aspect64
+
+# Smaller sizes stop discriminating: the calls either stay asynchronous or
+# block too briefly for a stable measurement.
+_BATCH = 512
+_SIZE = 4096
+
+# A call that holds the GIL measures around 0.09 of the reference rate, one
+# that releases it around 0.3. The threshold sits between the two.
+_MIN_RATIO = 0.18
+
+_BACKLOG = 2  # queued transforms, so the measured call has something to wait on
+
+_TRIALS = 5  # samples averaged per measurement
+
+
+class _Ticker:
+    """Counts how often a competing Python thread gets scheduled."""
+
+    def __enter__(self):
+        self.ticks = 0
+        self._stop = False
+        self._thread = threading.Thread(target=self._run, daemon=True)
+        self._thread.start()
+        return self
+
+    def _run(self):
+        while not self._stop:
+            self.ticks += 1
+            time.sleep(0)
+
+    def __exit__(self, *exc):
+        self._stop = True
+        self._thread.join(timeout=5)
+        return False
+
+
+@pytest.mark.skipif(not has_support_aspect64(), reason="requires fp64 support")
+class TestFftReleasesGil:
+    @pytest.fixture(autouse=True)
+    def _switch_interval(self):
+        # Stop CPython handing the GIL over on its own timer.
+        previous = sys.getswitchinterval()
+        sys.setswitchinterval(0.0005)
+        yield
+        sys.setswitchinterval(previous)
+
+    def _assert_releases_gil(self, name, a, fn):
+        queue = a.sycl_queue
+
+        def measure(ticker, func):
+            total_ticks = 0
+            total_s = 0.0
+            for _ in range(_TRIALS):
+                for _ in range(_BACKLOG):
+                    dpnp.fft.fft(a)
+                ticker.ticks = 0
+                start = time.perf_counter()
+                func()
+                total_s += time.perf_counter() - start
+                total_ticks += ticker.ticks
+                queue.wait()
+            # ticks per millisecond, and the average duration of a single call
+            return total_ticks / max(1000 * total_s, 1e-3), total_s / _TRIALS
+
+        fn()  # warm up JIT
+        queue.wait()
+
+        with _Ticker() as ticker:
+            measured, duration = measure(ticker, fn)
+            # time.sleep() is nogil, so blocking in it for the same amount of
+            # time gives the best tick rate the machine can produce
+            reference, _ = measure(ticker, lambda: time.sleep(duration))
+
+        assert reference > 0, "reference measurement produced no ticks"
+        ratio = measured / reference
+        assert ratio >= _MIN_RATIO, (
+            f"{name} holds the GIL while blocking: {measured:.2f} ticks/ms vs "
+            f"{reference:.2f} for a nogil sleep of the same duration (ratio "
+            f"{ratio:.3f}, need >= {_MIN_RATIO}). The oneMKL call needs "
+            f"py::gil_scoped_release."
+        )
+
+    @pytest.mark.slow
+    def test_fft_out_of_place(self):
+        # a complex input is passed to oneMKL as is, so the transform is
+        # computed out-of-place
+        a = dpnp.ones((_BATCH, _SIZE), dtype="c16")
+        self._assert_releases_gil(
+            "compute_fft_out_of_place", a, lambda: dpnp.fft.fft(a)
+        )
+
+    @pytest.mark.slow
+    def test_fft_in_place(self):
+        # a real input is copied to a complex array first, which allows the
+        # transform to be computed in-place
+        a = dpnp.ones((_BATCH, _SIZE), dtype="f8")
+        self._assert_releases_gil(
+            "compute_fft_in_place", a, lambda: dpnp.fft.fft(a)
+        )