Replace g_launch_mutex with GlobalLock

Improve safety via static typing techniques.

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

cext/CMakeLists.txt

@@ -126,7 +126,7 @@ endfunction()
 
 
 add_test_executable(test_stream_buffer
-    test/test_stream_buffer.cpp cuda_loader.cpp cuda_helper.cpp memory.cpp)
+    test/test_stream_buffer.cpp cuda_loader.cpp cuda_helper.cpp memory.cpp py.cpp)
 
 add_test_executable(test_hash_map test/test_hash_map.cpp memory.cpp)
 

cext/compiled_host.cpp

@@ -188,15 +188,16 @@ int CompiledHostProgram_init(PyObject* self, PyObject* args, PyObject* kwargs) {
 
 PyObject* invoke_host_entry(
         CompiledHostProgram& program,
-        void** arguments) {
+        void** arguments,
+        GlobalLock& lock) {
     int32_t result = program.executable.entry(arguments, &program.runtime);
     if (result < 0) {
         if (PyErr_Occurred()) return nullptr;
         raise(PyExc_RuntimeError, "compiled host code failed with status ", result);
         return nullptr;
     }
     if (result != CUDA_SUCCESS) {
-        Result<const DriverApi*> driver_result = get_driver_api();
+        Result<const DriverApi*> driver_result = get_driver_api(lock);
         if (!driver_result.is_ok()) return nullptr;
         const DriverApi* driver = *driver_result;
         raise(PyExc_RuntimeError, "cuda error occurred: ",
@@ -212,6 +213,7 @@ PyObject* CompiledHostProgram_invoke(PyObject* self, PyObject* argument_addresse
         raise(PyExc_TypeError, "compiled host argument addresses must be a tuple");
         return nullptr;
     }
+    GlobalLock lock;
     Py_ssize_t count = PyTuple_GET_SIZE(argument_addresses);
     Vec<void*> arguments;
     arguments.reserve(count);
@@ -221,7 +223,7 @@ PyObject* CompiledHostProgram_invoke(PyObject* self, PyObject* argument_addresse
         if (PyErr_Occurred()) return nullptr;
         arguments.push_back(address);
     }
-    return compiled_host_program_invoke(self, arguments.data());
+    return compiled_host_program_invoke(self, arguments.data(), lock);
 }
 
 
@@ -286,7 +288,8 @@ bool compiled_host_program_check(PyObject* object) {
 }
 
 
-PyObject* compiled_host_program_invoke(PyObject* program_object, void** arguments) {
+PyObject* compiled_host_program_invoke(PyObject* program_object, void** arguments,
+                                       GlobalLock& lock) {
     if (!compiled_host_program_check(program_object)) {
         raise(
                 PyExc_TypeError,
@@ -295,7 +298,7 @@ PyObject* compiled_host_program_invoke(PyObject* program_object, void** argument
         return nullptr;
     }
     CompiledHostProgram& program = py_unwrap<CompiledHostProgram>(program_object);
-    return invoke_host_entry(program, arguments);
+    return invoke_host_entry(program, arguments, lock);
 }
 
 

cext/compiled_host.h

@@ -8,5 +8,5 @@
 
 
 bool compiled_host_program_check(PyObject* object);
-PyObject* compiled_host_program_invoke(PyObject* program, void** arguments);
+PyObject* compiled_host_program_invoke(PyObject* program, void** arguments, GlobalLock& lock);
 Status compiled_host_init(PyObject* module);

cext/cuda_helper.cpp

@@ -37,7 +37,8 @@ PyObject* get_max_grid_size(PyObject *self, PyObject *args) {
     if (!PyArg_ParseTuple(args, "i", &device_id))
         return nullptr;
 
-    Result<const DriverApi*> driver = get_driver_api();
+    GlobalLock lock;
+    Result<const DriverApi*> driver = get_driver_api(lock);
     if (!driver.is_ok()) return nullptr;
 
     CUdevice dev;
@@ -123,7 +124,8 @@ PyObject* get_compute_capability(PyObject *self, PyObject *args) {
     int device_id = 0;
     if (!PyArg_ParseTuple(args, "|i", &device_id)) return nullptr;
 
-    Result<const DriverApi*> driver_result = get_driver_api();
+    GlobalLock lock;
+    Result<const DriverApi*> driver_result = get_driver_api(lock);
     if (!driver_result.is_ok()) return nullptr;
 
     Result<ComputeCapability> computeCapability =
@@ -135,7 +137,8 @@ PyObject* get_compute_capability(PyObject *self, PyObject *args) {
 PyObject* get_driver_version(PyObject *self, PyObject *Py_UNUSED(ignored)) {
     int major, minor;
 
-    Result<const DriverApi*> driver_result = get_driver_api();
+    GlobalLock lock;
+    Result<const DriverApi*> driver_result = get_driver_api(lock);
     if (!driver_result.is_ok()) return nullptr;
     const DriverApi* d = *driver_result;
 
@@ -152,7 +155,8 @@ PyObject* get_driver_version(PyObject *self, PyObject *Py_UNUSED(ignored)) {
 // ========== Context helpers ==========
 
 PyObject* synchronize_context(PyObject* self, PyObject* Py_UNUSED(ignored)) {
-    Result<const DriverApi*> driver_result = get_driver_api();
+    GlobalLock lock;
+    Result<const DriverApi*> driver_result = get_driver_api(lock);
     if (!driver_result.is_ok()) return nullptr;
     const DriverApi* d = *driver_result;
 
@@ -167,7 +171,8 @@ PyObject* synchronize_context(PyObject* self, PyObject* Py_UNUSED(ignored)) {
 // ========== Stream helpers ==========
 
 PyObject* create_stream(PyObject* self, PyObject* Py_UNUSED(ignored)) {
-    Result<const DriverApi*> driver_result = get_driver_api();
+    GlobalLock lock;
+    Result<const DriverApi*> driver_result = get_driver_api(lock);
     if (!driver_result.is_ok()) return nullptr;
     const DriverApi* d = *driver_result;
 
@@ -181,10 +186,11 @@ PyObject* create_stream(PyObject* self, PyObject* Py_UNUSED(ignored)) {
 }
 
 PyObject* destroy_stream(PyObject* self, PyObject* arg) {
+    GlobalLock lock;
     CUstream stream = static_cast<CUstream>(PyLong_AsVoidPtr(arg));
     if (PyErr_Occurred()) return nullptr;
 
-    Result<const DriverApi*> driver_result = get_driver_api();
+    Result<const DriverApi*> driver_result = get_driver_api(lock);
     if (!driver_result.is_ok()) return nullptr;
     const DriverApi* d = *driver_result;
 
@@ -226,15 +232,14 @@ static CUresult shim_cuLaunchKernelEx(
 }
 
 static PyObject* spy_on_cuLaunchKernel_begin(PyObject* self, PyObject* arg) {
-#ifdef Py_GIL_DISABLED
-    PyCriticalSectionGuard guard(&g_spy_mutex);
-#endif
+    GlobalLock lock;
+
     if (g_real_cuLaunchKernelEx) {
         raise(PyExc_RuntimeError, "Already spying");
         return nullptr;
     }
 
-    Result<const DriverApi*> driver_result = get_driver_api();
+    Result<const DriverApi*> driver_result = get_driver_api(lock);
     if (!driver_result.is_ok()) return nullptr;
 
     DriverApi* api = const_cast<DriverApi*>(*driver_result);
@@ -245,15 +250,14 @@ 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
+    GlobalLock lock;
+
     if (!g_real_cuLaunchKernelEx) {
         raise(PyExc_RuntimeError, "Not spying");
         return nullptr;
     }
 
-    Result<const DriverApi*> driver_result = get_driver_api();
+    Result<const DriverApi*> driver_result = get_driver_api(lock);
     if (!driver_result.is_ok()) return nullptr;
 
     DriverApi* api = const_cast<DriverApi*>(*driver_result);

cext/cuda_loader.cpp

@@ -73,7 +73,7 @@ static Result<cuGetProcAddress_v2_t> get_cuGetProcAddress_from_python() {
 
 static constexpr int MIN_DRIVER_VERSION = 13000;
 
-Result<const DriverApi*> get_driver_api() {
+Result<const DriverApi*> get_driver_api(GlobalLock& lock) {
     static bool initialized;
     static DriverApi instance;
     if (!initialized) {

cext/cuda_loader.h

@@ -82,7 +82,7 @@ typedef CUresult (*cuGetProcAddress_v2_t)
 
 Status driver_api_init(DriverApi* driver_api, cuGetProcAddress_v2_t _cuGetProcAddress);
 
-Result<const DriverApi*> get_driver_api();
+Result<const DriverApi*> get_driver_api(GlobalLock& lock);
 
 
 class CudaContextGuard {

cext/py.cpp

@@ -47,3 +47,7 @@ void log_python_error(const char* filename, int line, const char* level, SavedEx
     PyErr_SetExcInfo(old_excinfo_type, old_excinfo_value, old_excinfo_tb);
 }
 
+#ifdef Py_GIL_DISABLED
+PyMutex GlobalLock::mutex_ = {0};
+#endif
+

cext/py.h

@@ -601,3 +601,75 @@ class PyCriticalSectionGuard {
         PyCriticalSection _py_cs;
 };
 #endif
+
+
+// In a GIL build, asserts that we are holding the GIL.
+// In a free-threaded build, enters the global critical section.
+//
+// Passing a reference to an object of this class to a function serves as a "proof"
+// that the lock is being held.
+//
+// For example:
+//
+//     // Helper function that mutates the global state in a thread-unsafe way.
+//     // Takes a `GlobalLock&` reference to indicate that either the GIL or the global
+//     // critical section must be held.
+//     static int next_number(GlobalLock&) {
+//         static int counter = 0;
+//         return counter++;
+//     }
+//
+//     // Method implementation
+//     static PyObject* foo(PyObject* self, PyObject* args) {
+//         // In a GIL build, we know we're holding the GIL since this is a method, and
+//         // thus it's safe to instantiate a GlobalLock.
+//         //
+//         // In a free-threaded build, this will enter the global critical section
+//         // to emulate the GIL.
+//         GlobalLock lock;
+//
+//         // Call the thread-unsafe helper.
+//         int number = next_number(lock);
+//
+//         // ...
+//     }
+class GlobalLock {
+public:
+    GlobalLock()
+#ifdef Py_GIL_DISABLED
+        : guard_(&mutex_)
+#endif
+    {
+#ifndef Py_GIL_DISABLED
+        // Will crash with a fatal error if we aren't holding the GIL
+        PyThreadState_Get();
+#endif // Py_GIL_DISABLED
+    }
+
+    // Use this constructor when you are holding a GILGuard
+    // but you don't necessarily have an attached thread state.
+    explicit GlobalLock(GILGuard& gil_guard)
+#ifdef Py_GIL_DISABLED
+        : guard_(&mutex_)
+#endif
+    {}
+
+    GlobalLock(const GlobalLock&) = delete;
+    void operator==(const GlobalLock&) = delete;
+
+private:
+#ifdef Py_GIL_DISABLED
+    PyCriticalSectionGuard guard_;
+    static PyMutex mutex_;
+#endif
+};
+
+
+template <typename T>
+class ProtectedByGlobalLock {
+    static_assert(std::is_trivially_destructible_v<T>);
+    T object_;
+public:
+    T& get(GlobalLock&) { return object_; }
+};
+