Support py314 free-threading

gunnersdeng · gunnersdeng · commit f31c7629d0f6 · 2026-05-20T10:31:31.000-07:00
Signed-off-by: Ziheng Deng &lt;zihengd@nvidia.com&gt;
diff --git a/README.md b/README.md
@@ -149,10 +149,17 @@ Running Tests
 -------------
 cuTile uses the [pytest](https://pytest.org) framework for testing.
 Tests have extra dependencies, such as PyTorch, which can be installed with
+
+For Python non-free-threading build:
 ```
 pip install -r test/requirements.txt
 ```
 
+Or for Python free-threading build:
+```
+pip install -r test/requirements-ft.txt
+```
+
 The tests are located in the [test/](test/) directory. To run a specific test file,
 for example `test_copy.py`, use the following command:
 ```
diff --git a/cext/CMakeLists.txt b/cext/CMakeLists.txt
@@ -12,8 +12,23 @@ if (MSVC)
     set(cext_compile_flags /GR- /GS- /EHs- -DNDEBUG)
     # /NODEFAULTLIB and /NOENTRY to disable dependency on runtime lib msvcrt.dll
     set(cext_link_flags /NODEFAULTLIB /NOENTRY)
+    # set Py_GIL_DISABLED macro, at it is not defined automatically on Windows
+    execute_process(
+        COMMAND ${Python_EXECUTABLE} -c
+            "import sysconfig; print(sysconfig.get_config_var('Py_GIL_DISABLED') or 0)"
+        OUTPUT_VARIABLE PY_GIL_DISABLED
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+    )
+    if (PY_GIL_DISABLED EQUAL 1)
+        list(APPEND cext_compile_flags -DPy_GIL_DISABLED=1)
+    endif()
 else()
     set(cext_compile_flags -fno-exceptions -fno-rtti -fPIC -fvisibility=hidden)
+    # For arm64 add flag to disable calls to out-of-line helpers for atomic operations 
+    # to avoid undefined reference error during compilation
+    if (CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$")
+        list(APPEND cext_compile_flags -mno-outline-atomics)
+    endif()
     set(cext_link_flags ${cext_compile_flags})
     set(nostdlib_flags -nostdlib -DNDEBUG -fno-builtin)
 endif()
diff --git a/cext/cuda_helper.cpp b/cext/cuda_helper.cpp
@@ -143,15 +143,21 @@ PyObject* destroy_stream(PyObject* self, PyObject* arg) {
     Py_RETURN_NONE;
 }
 
-static decltype(cuLaunchKernelEx)* g_real_cuLaunchKernelEx;
-static PyObject* g_cuLaunchKernelEx_spy_callback;
+static decltype(cuLaunchKernelEx)* g_real_cuLaunchKernelEx; // Protected by the GIL or g_spy_mutex
+static PyObject* g_cuLaunchKernelEx_spy_callback; // Protected by the GIL or g_spy_mutex
+
+#ifdef Py_GIL_DISABLED
+static PyMutex g_spy_mutex = {0};
+#endif
 
 static CUresult shim_cuLaunchKernelEx(
         const CUlaunchConfig *config,
         CUfunction f,
         void** kernelParams,
         void** extra) {
-
+#ifdef Py_GIL_DISABLED
+    PyCriticalSectionGuard guard(&g_spy_mutex);
+#endif
     PyPtr res = steal(PyObject_CallFunction(
             g_cuLaunchKernelEx_spy_callback,
             "(K III III I K)",
@@ -167,6 +173,9 @@ static CUresult shim_cuLaunchKernelEx(
 }
 
 static PyObject* spy_on_cuLaunchKernel_begin(PyObject* self, PyObject* arg) {
+#ifdef Py_GIL_DISABLED
+    PyCriticalSectionGuard guard(&g_spy_mutex);
+#endif
     if (g_real_cuLaunchKernelEx)
         return PyErr_Format(PyExc_RuntimeError, "Already spying");
 
@@ -181,6 +190,9 @@ static PyObject* spy_on_cuLaunchKernel_begin(PyObject* self, PyObject* arg) {
 }
 
 static PyObject* spy_on_cuLaunchKernel_end(PyObject* self, PyObject* arg) {
+#ifdef Py_GIL_DISABLED
+    PyCriticalSectionGuard guard(&g_spy_mutex);
+#endif
     if (!g_real_cuLaunchKernelEx)
         return PyErr_Format(PyExc_RuntimeError, "Not spying");
 
diff --git a/cext/cuda_loader.cpp b/cext/cuda_loader.cpp
@@ -77,9 +77,16 @@ static Status cuda_loader_init(DriverApi& driver_api) {
 
 static constexpr int MIN_DRIVER_VERSION = 13000;
 
+#ifdef Py_GIL_DISABLED
+static PyMutex g_driver_api_mutex = {0};
+#endif
+
 Result<const DriverApi*> get_driver_api() {
     static bool initialized;
     static DriverApi instance;
+#ifdef Py_GIL_DISABLED
+    PyCriticalSectionGuard guard(&g_driver_api_mutex);
+#endif
     if (!initialized) {
         if (!cuda_loader_init(instance))
             return ErrorRaised;
diff --git a/cext/module.cpp b/cext/module.cpp
@@ -36,6 +36,11 @@ PyMODINIT_FUNC PyInit__cext() {
     PyPtr m = steal(PyModule_Create(&module_def));
     if (!m) return nullptr;
 
+#ifdef Py_GIL_DISABLED
+    if (PyUnstable_Module_SetGIL(m.get(), Py_MOD_GIL_NOT_USED) != 0 )
+        return nullptr;
+#endif
+
     if (!tile_kernel_init(m.get()))
         return nullptr;
 
diff --git a/cext/py.h b/cext/py.h
@@ -6,6 +6,7 @@
 
 #pragma once
 
+#include "check.h"
 #include "ref_ptr.h"
 #include <Python.h>
 #include <optional>
@@ -275,3 +276,22 @@ static inline PyPtr try_import(const char* modname, SavedException* exc = nullpt
     if (!ret && exc) *exc = save_raised_exception();
     return ret;
 }
+
+#ifdef Py_GIL_DISABLED
+class PyCriticalSectionGuard {
+    public:
+        explicit PyCriticalSectionGuard(PyMutex* mutex) {
+            CHECK(mutex);
+            PyCriticalSection_BeginMutex(&_py_cs, mutex);
+        }
+
+        ~PyCriticalSectionGuard() {
+            PyCriticalSection_End(&_py_cs);
+        }
+
+        PyCriticalSectionGuard(const PyCriticalSectionGuard&) = delete;
+        void operator=(const PyCriticalSectionGuard&) = delete;
+    private:
+        PyCriticalSection _py_cs;
+};
+#endif
diff --git a/cext/tile_kernel.cpp b/cext/tile_kernel.cpp
@@ -377,7 +377,7 @@ static Word* push_single_word_cuarg(LaunchHelper& helper, Word word) {
     return ptr;
 }
 
-static LaunchHelper* g_helper_freelist;  // protected by the GIL
+static LaunchHelper* g_helper_freelist;  // protected by the GIL or g_launch_mutex
 
 namespace { struct LaunchHelperDeleter {
     void operator() (LaunchHelper* helper) const {
@@ -386,6 +386,10 @@ namespace { struct LaunchHelperDeleter {
     }
 }; }
 
+#ifdef Py_GIL_DISABLED
+static PyMutex g_launch_mutex = {0};
+#endif
+
 using LaunchHelperPtr = std::unique_ptr<LaunchHelper, LaunchHelperDeleter>;
 
 
@@ -1489,6 +1493,9 @@ struct CompareKey <Vec<PyTypeObject*>, Vec<PyPtr>> {
 namespace { struct TileContext {
     PyPtr config;
     PyPtr autotune_cache;
+#ifdef Py_GIL_DISABLED
+    PyMutex accessor_mutex = {0};
+#endif
 
     static PyTypeObject pytype;
 }; }
@@ -1913,7 +1920,7 @@ static Result<CUstream> parse_stream(PyObject* py_stream) {
 
 using StreamBufferPoolMap = HashMap<unsigned long long, StreamBufferPool*>;
 
-// Protected by GIL.
+// Protected by GIL or g_launch_mutex.
 // We have no reliable way to detect when a context is destroyed, so we never clean these up.
 static StreamBufferPoolMap* g_stream_buffer_pool_by_ctx_id;
 
@@ -2413,16 +2420,27 @@ static int TileContext_init(PyObject* self, PyObject* args, PyObject* kwargs) {
 
 
 static PyObject * TileContext_get_config(PyObject* self, void *closure) {
-    return Py_NewRef(py_unwrap<TileContext>(self).config.get());
+    TileContext& context = py_unwrap<TileContext>(self);
+#ifdef Py_GIL_DISABLED
+    PyCriticalSectionGuard guard(&context.accessor_mutex);
+#endif
+    return Py_NewRef(context.config.get());
 }
 
 
 static PyObject * TileContext_get_autotune_cache(PyObject* self, void *closure) {
-    return Py_NewRef(py_unwrap<TileContext>(self).autotune_cache.get());
+    TileContext& context = py_unwrap<TileContext>(self);
+#ifdef Py_GIL_DISABLED
+    PyCriticalSectionGuard guard(&context.accessor_mutex);
+#endif
+    return Py_NewRef(context.autotune_cache.get());
 }
 
 static int TileContext_set_autotune_cache(PyObject* self, PyObject* value, void* closure) {
     TileContext& context = py_unwrap<TileContext>(self);
+#ifdef Py_GIL_DISABLED
+    PyCriticalSectionGuard guard(&context.accessor_mutex);
+#endif
 
     // `del ctx.autotune_cache` → set back to None
     if (value == nullptr) {
@@ -2490,6 +2508,9 @@ PyTypeObject TileDispatcher::pytype = {
 };
 
 static PyObject* get_parameter_constraints_from_pyargs(PyObject* self, PyObject* args) {
+#ifdef Py_GIL_DISABLED
+    PyCriticalSectionGuard guard(&g_launch_mutex);
+#endif
     PyObject* dispatcher_pyobj = nullptr;
     PyObject* pyargs = nullptr;
     PyObject* cconv = nullptr;
@@ -2683,6 +2704,9 @@ static Status parse_launch_args(PyObject* const* args, Py_ssize_t nargs, const c
 static PyObject* launch_impl(PyObject* const* args, Py_ssize_t nargs,
                              PyObject* kwargs, const char* signature, bool with_block
                              ) {
+#ifdef Py_GIL_DISABLED
+    PyCriticalSectionGuard guard(&g_launch_mutex);
+#endif
     LaunchArgs launch_args;
     if (!parse_launch_args(args, nargs, signature, with_block, &launch_args))
         return nullptr;
@@ -2723,6 +2747,9 @@ static PyObject *launch_extended(PyObject *, PyObject *const *args,
 #define BENCHMARK_SIGNATURE "_benchmark(stream, grid, kernel, pyargs_tuples, /)"
 
 static PyObject* cuda_tile_benchmark(PyObject* mod, PyObject* const* args, Py_ssize_t nargs) {
+#ifdef Py_GIL_DISABLED
+    PyCriticalSectionGuard guard(&g_launch_mutex);
+#endif
     LaunchArgs launch_args;
     if (!parse_launch_args(args, nargs, BENCHMARK_SIGNATURE, false, &launch_args))
         return nullptr;
diff --git a/changelog.d/support-py314.md b/changelog.d/support-py314.md
@@ -0,0 +1,2 @@
+- Add Python3.14 support
+- Add Python3.14 free threading support, safe to launch kernels via different threads with GIL disabled. 
diff --git a/setup.py b/setup.py
@@ -42,6 +42,7 @@ def _cmake(self, build_dir: str, build_type: str, dlpack_path: str):
         cmake_cmd = ["cmake", "-B", build_dir, project_root,
                      f"-DDLPACK_PATH={dlpack_path}",
                      f"-DCMAKE_BUILD_TYPE={build_type}",
+                     f"-DPython_EXECUTABLE={sys.executable}",
                      "-DCMAKE_POLICY_VERSION_MINIMUM=3.5"]
         if self.disable_internal:
             cmake_cmd.append("-DDISABLE_INTERNAL=1")
diff --git a/test/requirements-ft.txt b/test/requirements-ft.txt
@@ -0,0 +1,21 @@
+# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+setuptools==80.9.0
+pytest==9.0.1
+pytest-benchmark==5.1.0
+pytest-cov==6.2.1
+pytest-env==1.2.0
+flake8==7.3.0
+pytest-cov==6.2.1
+
+numpy==2.2.0; python_version < "3.11"
+numpy==2.4.4; python_version >= "3.11"
+
+jax[cuda13]==0.10.0; sys_platform == 'linux' and python_version >= "3.11"
+flatbuffers==25.12.19; sys_platform == 'linux' and python_version >= "3.11"
+
+
+--extra-index-url https://download.pytorch.org/whl/cu130
+torch==2.10.0
diff --git a/test/requirements.txt b/test/requirements.txt
@@ -2,16 +2,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-setuptools==80.9.0
-pytest==9.0.1
-pytest-benchmark==5.1.0
-pytest-cov==6.2.1
-pytest-env==1.2.0
-flake8==7.3.0
-pytest-cov==6.2.1
+# Maintain a separate file for free-threading version. This can be removed when ABI environment markers are supported.(PEP 780)
+-r requirements-ft.txt
 
 numba-cuda[cu13]==0.30.0
 cupy-cuda13x==14.0.1
-
---extra-index-url https://download.pytorch.org/whl/cu130
-torch==2.10.0
diff --git a/test/test_concurrent_launch.py b/test/test_concurrent_launch.py
@@ -0,0 +1,91 @@
+# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+import subprocess
+import threading
+import torch
+import cuda.tile as ct
+import cuda.tile._cext as _cext
+import traceback
+import pytest
+
+
+@ct.kernel
+def vector_add(x, y, z, TILE: ct.Constant[int]):
+    bid = ct.bid(0)
+    tx = ct.load(x, index=(bid,), shape=(TILE,))
+    ty = ct.load(y, index=(bid,), shape=(TILE,))
+    ct.store(z, index=(bid,), tile=tx + ty)
+
+
+_TILES = (32, 64, 128, 256)
+
+_LAUNCH_FUNCS = {
+    "ct.launch": ct.launch,
+    "_cext.launch": _cext.launch,
+    "_cext._benchmark": _cext._benchmark,
+}
+
+
+def _worker(thread_id, n_iters, x, y, expected, errors, launch_func):
+    try:
+        for i in range(n_iters):
+            tile = _TILES[(thread_id + i) % len(_TILES)]
+            grid = (x.shape[0] // tile, 1, 1)
+            z = torch.zeros_like(x)
+            launch_func(torch.cuda.current_stream(), grid, vector_add, (x, y, z, tile))
+            torch.cuda.synchronize()
+            if not torch.equal(z, expected):
+                errors.append(f"thread={thread_id} iter={i} tile={tile}:"
+                              f" result mismatch: {z} != {expected}")
+                return
+    except Exception:
+        errors.append(f"thread={thread_id} exception: {traceback.format_exc()}")
+
+
+def _run_concurrent_launch(launch_func_name):
+    launch_func = _LAUNCH_FUNCS[launch_func_name]
+    N_THREADS = 8
+    N_ITERS = 800
+    N_ELEMS = 8192
+
+    x = torch.randint(0, 100, (N_ELEMS,), device='cuda', dtype=torch.int32)
+    y = torch.randint(0, 100, (N_ELEMS,), device='cuda', dtype=torch.int32)
+    expected = x + y
+
+    errors = []
+    threads = []
+    for t in range(N_THREADS):
+        t = threading.Thread(target=_worker, args=(t, N_ITERS, x, y, expected, errors, launch_func))
+        threads.append(t)
+        t.start()
+
+    for t in threads:
+        t.join()
+
+    if errors:
+        print("concurrent launch errors:\n  " + "\n  ".join(errors), file=sys.stderr)
+        sys.exit(1)
+
+
+@pytest.mark.parametrize("launch_func_name", list(_LAUNCH_FUNCS))
+def test_concurrent_launch(launch_func_name):
+    # Free-threading mode may result in segfaults, run in a subprocess to isolate the test.
+    result = subprocess.run(
+        [sys.executable, __file__, "run_concurrent_launch", launch_func_name], capture_output=True,
+    )
+    if result.returncode != 0:
+        print("--- Captured stdout ---")
+        print(result.stdout.decode(errors='replace'))
+        print("--- End Captured stdout ---\n--- Captured stderr ---")
+        print(result.stderr.decode(errors='replace'))
+        print("--- End Captured stderr ---")
+        pytest.fail(f"concurrent launch subprocess failed with return code {result.returncode}")
+
+
+if __name__ == "__main__":
+    cmd = sys.argv[1]
+    if cmd == "run_concurrent_launch":
+        _run_concurrent_launch(sys.argv[2])

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+- Add Python3.14 support`
	`2`	`+- Add Python3.14 free threading support, safe to launch kernels via different threads with GIL disabled.`