Make jax_cuda kernels deterministic and allow using multiple devices with the FFI kernel

CMakeLists.txt

@@ -63,7 +63,7 @@ if(CMAKE_CUDA_COMPILER)
                                             CUDA_SEPARABLE_COMPILATION ON)
     set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -rdc=true")
     set(CMAKE_CUDA_ARCHITECTURES
-        "70;80;89"
+        "70;75;80;86;89;90"
         CACHE STRING "List of CUDA compute capabilities to build cuDecomp for.")
     message(STATUS "CUDA_ARCHITECTURES: ${CMAKE_CUDA_ARCHITECTURES}")
     set_target_properties(_s2fft PROPERTIES CUDA_ARCHITECTURES

lib/include/cudadeviceguard.hpp

@@ -0,0 +1,40 @@
+#ifndef CUDADEVICEGUARD_HPP
+#define CUDADEVICEGUARD_HPP
+
+#include <cuda_runtime.h>
+
+/**
+ * @brief RAII selection of the device the current scope operates on.
+ *
+ * XLA may run FFI handlers from host threads whose current device is not the device owning
+ * the buffers (single-process multi-GPU, astro-informatics/s2fft#389). cuFFT plans and
+ * streams bind to the device that is current when they are created, so the execution device
+ * must be selected before any plan or stream is touched. The previous device is restored on
+ * scope exit; a negative ordinal is a no-op (host-side / device-agnostic callers).
+ */
+class CudaDeviceScope {
+public:
+    explicit CudaDeviceScope(int device) : m_restore(-1) {
+        if (device >= 0) {
+            int current = -1;
+            if (cudaGetDevice(&current) == cudaSuccess && current != device) {
+                cudaSetDevice(device);
+                m_restore = current;
+            }
+        }
+    }
+
+    ~CudaDeviceScope() {
+        if (m_restore >= 0) {
+            cudaSetDevice(m_restore);
+        }
+    }
+
+    CudaDeviceScope(const CudaDeviceScope &) = delete;
+    CudaDeviceScope &operator=(const CudaDeviceScope &) = delete;
+
+private:
+    int m_restore;
+};
+
+#endif  // CUDADEVICEGUARD_HPP

lib/include/cudastreamhandler.hpp

@@ -1,11 +1,14 @@
-
 /**
  * @file cudastreamhandler.hpp
- * @brief Singleton class for managing CUDA streams and events.
+ * @brief Singleton class for managing per-device CUDA streams and events.
+ *
+ * Streams and events are device-scoped CUDA objects: a stream created while device 0 is
+ * current cannot be used from another device. The pool is therefore keyed by device ordinal
+ * (single-process multi-GPU, astro-informatics/s2fft#389).
  *
- * This header provides a singleton implementation that encapsulates the creation,
- * management, and cleanup of CUDA streams and events. It offers functions to fork
- * streams, add new streams, and synchronize (join) streams with a given dependency.
+ * Events are one-shot dependencies: cudaStreamWaitEvent captures the event's recorded state
+ * at enqueue time, so each event is destroyed as soon as its wait is enqueued. A deferred
+ * cleanup thread (and its lock) is not needed and would contend with the transform handler.
  *
  * Usage example:
  * @code
@@ -14,15 +17,11 @@
  *   int main() {
  *       // Create a handler instance
  *       CudaStreamHandler handler;
- *
- *       // Fork 4 streams dependent on a given stream 'stream_main'
- *       handler.Fork(stream_main, 4);
- *
+ *       // Fork 4 streams on device 0, dependent on stream 'main'
+ *       auto forked = handler.Fork(0, main, 4);
  *       // Do work on the forked streams...
- *
- *       // Join the streams back to 'stream_main'
- *       handler.join(stream_main);
- *
+ *       // Join them back into 'main'
+ *       handler.join(main, forked);
  *       return 0;
  *   }
  * @endcode
@@ -34,114 +33,82 @@
 #define CUDASTREAMHANDLER_HPP
 
 #include <algorithm>
-#include <atomic>
-#include <cuda_runtime.h>
 #include <stdexcept>
-#include <thread>
+#include <unordered_map>
 #include <vector>
 
-// Singleton class managing CUDA streams and events
+#include <cuda_runtime.h>
+
+#include "cudadeviceguard.hpp"
+
+// Singleton class managing per-device CUDA streams and events
 class CudaStreamHandlerImpl {
 public:
     static CudaStreamHandlerImpl &instance() {
         static CudaStreamHandlerImpl instance;
         return instance;
     }
 
-    void AddStreams(int numStreams) {
-        if (numStreams > m_streams.size()) {
-            int streamsToAdd = numStreams - m_streams.size();
-            m_streams.resize(numStreams);
-            std::generate(m_streams.end() - streamsToAdd, m_streams.end(), []() {
+    /**
+     * @brief Forks N streams on `device`, each waiting on the current state of `dependentStream`.
+     *
+     * Returns exactly the N streams to use for this fork, so a later Fork with a smaller N can
+     * never hand out streams whose dependencies were set for a different batch count.
+     */
+    std::vector<cudaStream_t> Fork(int device, cudaStream_t dependentStream, int N) {
+        if (N <= 0) {
+            throw std::invalid_argument("Fork requires at least one stream.");
+        }
+        std::vector<cudaStream_t> &pool = m_streams_by_device[device];
+        if (pool.size() < static_cast<size_t>(N)) {
+            CudaDeviceScope guard(device);
+            while (pool.size() < static_cast<size_t>(N)) {
                 cudaStream_t stream;
-                cudaStreamCreate(&stream);
-                return stream;
-            });
+                if (cudaStreamCreate(&stream) != cudaSuccess) {
+                    throw std::runtime_error("Failed to create a CUDA stream.");
+                }
+                pool.push_back(stream);
+            }
         }
-    }
 
-    void join(cudaStream_t finalStream) {
-        std::for_each(m_streams.begin(), m_streams.end(), [this, finalStream](cudaStream_t stream) {
+        std::vector<cudaStream_t> forked(pool.end() - N, pool.end());
+        std::for_each(forked.begin(), forked.end(), [dependentStream](cudaStream_t stream) {
             cudaEvent_t event;
             cudaEventCreate(&event);
-            cudaEventRecord(event, stream);
-            cudaStreamWaitEvent(finalStream, event, 0);
-            m_events.push_back(event);
+            cudaEventRecord(event, dependentStream);
+            cudaStreamWaitEvent(stream, event, 0);  // Set the stream to wait on the event
+            cudaEventDestroy(event);
         });
-
-        if (!cleanup_thread.joinable()) {
-            stop_thread.store(false);
-            cleanup_thread = std::thread([this]() { this->AsyncEventCleanup(); });
-        }
+        return forked;
     }
 
-    // Fork function to add streams and set dependency on a given stream
-    void Fork(cudaStream_t dependentStream, int N) {
-        AddStreams(N);  // Add N streams
-
-        // Set dependency on the provided stream
-        std::for_each(m_streams.end() - N, m_streams.end(), [this, dependentStream](cudaStream_t stream) {
+    /**
+     * @brief Joins the forked streams back into `finalStream` (each waits on their completion).
+     */
+    void join(cudaStream_t finalStream, const std::vector<cudaStream_t> &forked) {
+        std::for_each(forked.begin(), forked.end(), [finalStream](cudaStream_t stream) {
             cudaEvent_t event;
             cudaEventCreate(&event);
-            cudaEventRecord(event, dependentStream);
-            cudaStreamWaitEvent(stream, event, 0);  // Set the stream to wait on the event
-            m_events.push_back(event);
+            cudaEventRecord(event, stream);
+            cudaStreamWaitEvent(finalStream, event, 0);
+            cudaEventDestroy(event);
         });
     }
 
-    auto getIterator() { return StreamIterator(m_streams.begin(), m_streams.end()); }
-
     ~CudaStreamHandlerImpl() {
-        stop_thread.store(true);
-        if (cleanup_thread.joinable()) {
-            cleanup_thread.join();
+        for (auto &entry : m_streams_by_device) {
+            std::for_each(entry.second.begin(), entry.second.end(), cudaStreamDestroy);
         }
-
-        std::for_each(m_streams.begin(), m_streams.end(), cudaStreamDestroy);
-        std::for_each(m_events.begin(), m_events.end(), cudaEventDestroy);
     }
 
-    // Custom Iterator class to iterate over streams
-    class StreamIterator {
-    public:
-        StreamIterator(std::vector<cudaStream_t>::iterator begin, std::vector<cudaStream_t>::iterator end)
-                : current(begin), end(end) {}
-
-        cudaStream_t next() {
-            if (current == end) {
-                throw std::out_of_range("No more streams.");
-            }
-            return *current++;
-        }
-
-        bool hasNext() const { return current != end; }
-
-    private:
-        std::vector<cudaStream_t>::iterator current;
-        std::vector<cudaStream_t>::iterator end;
-    };
-
 private:
-    CudaStreamHandlerImpl() : stop_thread(false) {}
+    CudaStreamHandlerImpl() = default;
     CudaStreamHandlerImpl(const CudaStreamHandlerImpl &) = delete;
     CudaStreamHandlerImpl &operator=(const CudaStreamHandlerImpl &) = delete;
 
-    void AsyncEventCleanup() {
-        while (!stop_thread.load()) {
-            std::for_each(m_events.begin(), m_events.end(), [this](cudaEvent_t &event) {
-                if (cudaEventQuery(event) == cudaSuccess) {
-                    cudaEventDestroy(event);
-                    event = nullptr;
-                }
-            });
-            std::this_thread::sleep_for(std::chrono::milliseconds(10));
-        }
-    }
-
-    std::vector<cudaStream_t> m_streams;
-    std::vector<cudaEvent_t> m_events;
-    std::thread cleanup_thread;
-    std::atomic<bool> stop_thread;
+    // Streams keyed by device ordinal: stream creation binds them to the device that is
+    // current at creation time.
+    std::unordered_map<int, std::vector<cudaStream_t>> m_streams_by_device;
 };
 
 // Public class for encapsulating the singleton operations
@@ -150,15 +117,12 @@ class CudaStreamHandler {
     CudaStreamHandler() = default;
     ~CudaStreamHandler() = default;
 
-    void AddStreams(int numStreams) { CudaStreamHandlerImpl::instance().AddStreams(numStreams); }
-
-    void join(cudaStream_t finalStream) { CudaStreamHandlerImpl::instance().join(finalStream); }
-
-    void Fork(cudaStream_t cudastream, int N) { CudaStreamHandlerImpl::instance().Fork(cudastream, N); }
+    std::vector<cudaStream_t> Fork(int device, cudaStream_t cudastream, int N) {
+        return CudaStreamHandlerImpl::instance().Fork(device, cudastream, N);
+    }
 
-    // Get the custom iterator for CUDA streams
-    CudaStreamHandlerImpl::StreamIterator getIterator() {
-        return CudaStreamHandlerImpl::instance().getIterator();
+    void join(cudaStream_t cudastream, const std::vector<cudaStream_t> &forked) {
+        CudaStreamHandlerImpl::instance().join(cudastream, forked);
     }
 };
 

lib/include/plan_cache.h

@@ -44,10 +44,15 @@ class PlanCache {
      * @tparam T The complex type (cufftComplex or cufftDoubleComplex) of the s2fftExec instance.
      * @param descriptor The s2fftDescriptor containing the parameters for the FFT.
      * @param executor A shared_ptr that will point to the retrieved or newly initialized s2fftExec instance.
+     * @param device_ordinal The CUDA device the executor must belong to. cuFFT plans bind to the
+     *        device that is current when they are created, so the cache is keyed per device
+     *        (single-process multi-GPU, astro-informatics/s2fft#389). A negative ordinal is a
+     *        host-side warm-up lookup (never used by the runtime paths).
      * @return HRESULT indicating success (S_OK if new, S_FALSE if from cache) or failure.
      */
     template <typename T>
-    HRESULT GetS2FFTExec(s2fftDescriptor &descriptor, std::shared_ptr<s2fftExec<T>> &executor);
+    HRESULT GetS2FFTExec(s2fftDescriptor &descriptor, std::shared_ptr<s2fftExec<T>> &executor,
+                         int32_t device_ordinal);
 
     /**
      * @brief Clears all cached s2fftExec instances.
@@ -68,12 +73,15 @@ class PlanCache {
 private:
     bool is_initialized = false;
 
-    // Unordered maps to store cached s2fftExec instances for double and single precision
-    std::unordered_map<s2fftDescriptor, std::shared_ptr<s2fftExec<cufftDoubleComplex>>,
-                       std::hash<s2fftDescriptor>, std::equal_to<>>
+    // Per-device caches of s2fftExec instances for double and single precision. cuFFT plans
+    // live in the CUDA context of the device they were created on, so a process-global,
+    // descriptor-only cache would hand device-0 plans to kernels running on other devices.
+    std::unordered_map<int32_t,
+                       std::unordered_map<s2fftDescriptor, std::shared_ptr<s2fftExec<cufftDoubleComplex>>,
+                                          std::hash<s2fftDescriptor>, std::equal_to<>>>
             m_Descriptors64;
-    std::unordered_map<s2fftDescriptor, std::shared_ptr<s2fftExec<cufftComplex>>, std::hash<s2fftDescriptor>,
-                       std::equal_to<>>
+    std::unordered_map<int32_t, std::unordered_map<s2fftDescriptor, std::shared_ptr<s2fftExec<cufftComplex>>,
+                                                   std::hash<s2fftDescriptor>, std::equal_to<>>>
             m_Descriptors32;
 
     /**