From cc92840e5787d942da28c207f08bb5ba1c0d1cfe Mon Sep 17 00:00:00 2001 From: Wassim Kabalan Date: Sun, 13 Sep 2026 16:36:43 +0200 Subject: [PATCH 1/4] Making jax_cuda deterministic for gradients --- s2fft/utils/healpix_ffts.py | 38 +++++++------- tests/test_healpix_ffts.py | 99 +++++++++++++++++++++++++++++++++++++ 2 files changed, 116 insertions(+), 21 deletions(-) diff --git a/s2fft/utils/healpix_ffts.py b/s2fft/utils/healpix_ffts.py index ec93d501f..07b1e6426 100644 --- a/s2fft/utils/healpix_ffts.py +++ b/s2fft/utils/healpix_ffts.py @@ -704,34 +704,30 @@ def _healpix_fft_cuda_lowering(ctx, f, *, L, nside, reality, fft_type, norm, adj The result of the FFI call. """ - # Step 1: Get the abstract evaluation results for the outputs. + # Step 1: Check if CUDA support is compiled in. + if not _s2fft.COMPILED_WITH_CUDA: + raise MissingCUDASupport() + + # Step 2: Get the abstract evaluation results for the outputs. (_, aval_out, _) = ctx.avals_out - # Step 2: Get lowering information (double precision, forward/backward, normalize). + # Step 3: Get lowering information (double precision, forward/backward, normalize). is_double, forward, normalize = _get_lowering_info(fft_type, norm, aval_out.dtype) - # Step 3: Select the appropriate FFI lowering function based on precision. - # We use operand_output_aliases={0: 0} to tell XLA that the input buffer (operand 0) - # can be reused for the output buffer (output 0). This allows XLA to perform the - # operation in-place if possible. Crucially, JAX manages this aliasing: if the input - # buffer is needed elsewhere (e.g. for backpropagation or if the user holds a reference), - # JAX will automatically copy the input to a new buffer before passing it to the - # kernel, ensuring that the original input is not corrupted. - # User can force the operation to be in-place by donating the input buffer, e.g. via healpix_fft_cuda(..., donate_argnums=(0,)). - # Even though this is an intermediate output it will instruct XLA that this is safe to reuse for the output of the primitive, - # and thus allows the kernel to run in-place without unnecessary copying. - # However, this is generally not needed and we can (generally) rely on XLA's optimizations to drop the input buffer if unused - # For more info .. check the XLA HLO of the operation + # Step 4: Select the appropriate FFI lowering function based on precision. + # NO operand-output alias here, deliberately. Declaring result 0 as an alias of operand 0 + # (in-place reuse) made the transform consume uninitialized memory whenever XLA's custom + # fusion pass materialized the aliased result as a fresh buffer without copying the + # operand into it, and even with a handler-side copy the alias changes XLA's buffer + # liveness unsoundly around fused custom calls. Result 0 is therefore always a plain + # output buffer, and the CUDA handler copies the operand into it before transforming + # in place (see lib/src/extensions.cc) -- sound under every XLA lowering path. if is_double: - ffi_lowered = jax.ffi.ffi_lowering( - "healpix_fft_cuda_c128", operand_output_aliases={0: 0} - ) + ffi_lowered = jax.ffi.ffi_lowering("healpix_fft_cuda_c128") else: - ffi_lowered = jax.ffi.ffi_lowering( - "healpix_fft_cuda_c64", operand_output_aliases={0: 0} - ) + ffi_lowered = jax.ffi.ffi_lowering("healpix_fft_cuda_c64") - # Step 4: Call the FFI lowering function with the context and parameters. + # Step 5: Call the FFI lowering function with the context and parameters. return ffi_lowered( ctx, f, diff --git a/tests/test_healpix_ffts.py b/tests/test_healpix_ffts.py index 110ffa098..3c6e24f6d 100644 --- a/tests/test_healpix_ffts.py +++ b/tests/test_healpix_ffts.py @@ -220,6 +220,105 @@ def healpix_inv_cuda(ftm): ) +@pytest.mark.skipif(not gpu_available, reason="GPU not available") +def test_healpix_map2alm_cuda_jit_roundtrip_survives_xla_fusion(): + """Regression for the CUDA FFI aliasing contract under jit. + + Two independent map2alm/alm2map roundtrips on traced operands of one parent array in a + single jit trigger XLA's custom-fusion pass. With the operand-output alias declared, the + fused custom call consumed an uninitialized result buffer (zeros on the compile run, stale + data on replay). The handler now always materializes its own aliased result, so the jitted + roundtrip must match eager to float64 precision. + """ + import jax_healpy as jhp + + nside, lmax = 16, 2 * 16 - 1 + npix = 12 * nside**2 + rng = np.random.default_rng(0) + maps = jnp.asarray(rng.normal(size=(2, npix))) + + def roundtrip(m): + flm = jhp.map2alm( + m, lmax=lmax, iter=0, pol=False, healpy_ordering=False, method="jax_cuda" + ) + return jnp.real( + jhp.alm2map( + flm, + nside=nside, + lmax=lmax, + pol=False, + healpy_ordering=False, + method="jax_cuda", + ) + ) + + eager = (roundtrip(maps[0]), roundtrip(maps[1])) + jitted = jax.jit(lambda m: (roundtrip(m[0]), roundtrip(m[1])))(maps) + again = jax.jit(lambda m: (roundtrip(m[0]), roundtrip(m[1])))(maps) + for i in range(2): + assert_allclose(jitted[i], eager[i], atol=1e-12, rtol=1e-12) + assert_allclose(again[i], jitted[i], atol=1e-12, rtol=1e-12) + + +@pytest.mark.skipif(not gpu_available, reason="GPU not available") +def test_healpix_map2alm_cuda_jit_gradient_matches_finite_differences(): + """The value_and_grad of a jitted CUDA-SHT likelihood must match finite differences. + + Guards the adjoint (transpose) path of the CUDA primitive: with the broken alias contract + the jitted value_and_grad returned a garbage potential and a 100%-wrong gradient. + """ + import jax_healpy as jhp + + nside, lmax = 16, 2 * 16 - 1 + npix = 12 * nside**2 + rng = np.random.default_rng(0) + m = jnp.asarray(rng.normal(size=npix)) + obs = jnp.real( + jhp.alm2map( + jhp.map2alm( + m * 0.5, + lmax=lmax, + iter=0, + pol=False, + healpy_ordering=False, + method="jax", + ), + nside=nside, + lmax=lmax, + pol=False, + healpy_ordering=False, + method="jax", + ) + ) + + def chi2(m): + flm = jhp.map2alm( + m, lmax=lmax, iter=0, pol=False, healpy_ordering=False, method="jax_cuda" + ) + back = jnp.real( + jhp.alm2map( + flm, + nside=nside, + lmax=lmax, + pol=False, + healpy_ordering=False, + method="jax_cuda", + ) + ) + return jnp.sum(0.5 * (back - obs) ** 2) + + vg = jax.jit(jax.value_and_grad(chi2)) + f0, grad = vg(m) + + key = jax.random.PRNGKey(7) + v = jax.random.normal(key, m.shape) + v = v / jnp.sqrt(jnp.vdot(v, v).real) + h = 1e-5 + fd = float((chi2(m + h * v) - chi2(m - h * v)) / (2 * h)) + gd = float(jnp.vdot(grad, v).real) + assert abs(gd - fd) / max(abs(fd), 1e-30) < 1e-4 + + def test_extension_module_not_available_raise_error(monkeypatch): monkeypatch.setattr(sys.modules["s2fft.utils.healpix_ffts"], "_s2fft", None) with pytest.raises(MissingExtensionModule): From 78b291c8fe82564866264412e0caa16b3f90e151 Mon Sep 17 00:00:00 2001 From: Wassim Kabalan Date: Sun, 13 Sep 2026 16:37:12 +0200 Subject: [PATCH 2/4] Implementing DeviceOrd capabilities for multi gpu map2alm --- lib/include/cudadeviceguard.hpp | 40 +++++++ lib/include/cudastreamhandler.hpp | 160 +++++++++++----------------- lib/include/plan_cache.h | 20 ++-- lib/src/extensions.cc | 168 ++++++++++++++++++++---------- lib/src/plan_cache.cc | 34 +++--- 5 files changed, 248 insertions(+), 174 deletions(-) create mode 100644 lib/include/cudadeviceguard.hpp diff --git a/lib/include/cudadeviceguard.hpp b/lib/include/cudadeviceguard.hpp new file mode 100644 index 000000000..726c46e86 --- /dev/null +++ b/lib/include/cudadeviceguard.hpp @@ -0,0 +1,40 @@ +#ifndef CUDADEVICEGUARD_HPP +#define CUDADEVICEGUARD_HPP + +#include + +/** + * @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(¤t) == 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 diff --git a/lib/include/cudastreamhandler.hpp b/lib/include/cudastreamhandler.hpp index f1b4ab4da..af4ce2109 100644 --- a/lib/include/cudastreamhandler.hpp +++ b/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,13 +33,15 @@ #define CUDASTREAMHANDLER_HPP #include -#include -#include #include -#include +#include #include -// Singleton class managing CUDA streams and events +#include + +#include "cudadeviceguard.hpp" + +// Singleton class managing per-device CUDA streams and events class CudaStreamHandlerImpl { public: static CudaStreamHandlerImpl &instance() { @@ -48,100 +49,66 @@ class CudaStreamHandlerImpl { 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 Fork(int device, cudaStream_t dependentStream, int N) { + if (N <= 0) { + throw std::invalid_argument("Fork requires at least one stream."); + } + std::vector &pool = m_streams_by_device[device]; + if (pool.size() < static_cast(N)) { + CudaDeviceScope guard(device); + while (pool.size() < static_cast(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 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 &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::iterator begin, std::vector::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::iterator current; - std::vector::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 m_streams; - std::vector m_events; - std::thread cleanup_thread; - std::atomic stop_thread; + // Streams keyed by device ordinal: stream creation binds them to the device that is + // current at creation time. + std::unordered_map> 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 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 &forked) { + CudaStreamHandlerImpl::instance().join(cudastream, forked); } }; diff --git a/lib/include/plan_cache.h b/lib/include/plan_cache.h index 9038cb76f..fdb44828a 100644 --- a/lib/include/plan_cache.h +++ b/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 - HRESULT GetS2FFTExec(s2fftDescriptor &descriptor, std::shared_ptr> &executor); + HRESULT GetS2FFTExec(s2fftDescriptor &descriptor, std::shared_ptr> &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>, - std::hash, 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>, + std::hash, std::equal_to<>>> m_Descriptors64; - std::unordered_map>, std::hash, - std::equal_to<>> + std::unordered_map>, + std::hash, std::equal_to<>>> m_Descriptors32; /** diff --git a/lib/src/extensions.cc b/lib/src/extensions.cc index b7183d624..f64786705 100644 --- a/lib/src/extensions.cc +++ b/lib/src/extensions.cc @@ -11,6 +11,7 @@ namespace nb = nanobind; #include "xla/ffi/api/c_api.h" #include "xla/ffi/api/ffi.h" #include "cuda_runtime.h" +#include "cudadeviceguard.hpp" #include "plan_cache.h" #include "s2fft_kernels.h" #include "s2fft.h" @@ -53,6 +54,32 @@ struct is_double : std::true_type {}; template constexpr bool is_double_v = is_double::value; +/** + * @brief Copies the operand into the aliased result buffer when XLA materialized them apart. + * + * The FFI lowering declares result 0 (`input_alias`) as an alias of operand 0, promising XLA + * that this handler writes it. The transform below READS `input_alias` in place, so the input + * data must be present in that buffer. When XLA passes the operand's own buffer the pointers + * match and nothing happens; when XLA materializes result 0 as a fresh buffer (its custom + * fusion pass does exactly that around slice+bitcast+custom-call), the copy makes the read + * sound instead of consuming uninitialized memory. + */ +template +ffi::Error copy_operand_to_alias(cudaStream_t stream, const ffi::Buffer& input, + ffi::Result> input_alias) { + if (input_alias->typed_data() == input.typed_data()) { + return ffi::Error::Success(); + } + cudaError_t err = cudaMemcpyAsync(input_alias->typed_data(), input.typed_data(), input.size_bytes(), + cudaMemcpyDeviceToDevice, stream); + if (err != cudaSuccess) { + return ffi::Error::Internal( + std::string("s2fft: failed to copy the operand into the aliased buffer: ") + + cudaGetErrorString(err)); + } + return ffi::Error::Success(); +} + /** * @brief Performs a forward HEALPix transform on a single element or batch. * @@ -63,7 +90,7 @@ constexpr bool is_double_v = is_double::value; * * @tparam T The XLA data type (F32, F64, etc). * @param stream CUDA stream to use. - * @param scratch ScratchAllocator for temporary device memory. + * @param device_ordinal CUDA device the buffers and plans live on. * @param input Input buffer containing HEALPix pixel-space data. * @param output Output buffer to store the FTM result. * @param workspace Output buffer for temporary workspace memory. @@ -71,35 +98,39 @@ constexpr bool is_double_v = is_double::value; * @return ffi::Error indicating success or failure. */ template -ffi::Error healpix_forward(cudaStream_t stream, ffi::Buffer input, ffi::Result> input_alias, - ffi::Result> output, ffi::Result> workspace, - s2fftDescriptor descriptor) { +ffi::Error healpix_forward(cudaStream_t stream, int32_t device_ordinal, ffi::Buffer input, + ffi::Result> input_alias, ffi::Result> output, + ffi::Result> workspace, s2fftDescriptor descriptor) { // Step 1: Determine the complex type based on the XLA data type. using fft_complex_type = fft_complex_t; const auto& dim_in = input.dimensions(); - // Step 1a: Get shift strategy from descriptor. - bool is_batched = (dim_in.size() == 2); + // Step 1a: Make the input data visible in the aliased result buffer (no-op when they alias). + if (ffi::Error err = copy_operand_to_alias(stream, input, input_alias); !err.success()) { + return err; + } // Step 2: Handle batched and non-batched cases separately. + bool is_batched = (dim_in.size() == 2); + + // Step 2a: Handle batched and non-batched cases separately. if (is_batched) { - // Step 2a: Batched case. + // Step 2b: Batched case. int batch_count = dim_in[0]; - // Step 2b: Compute offsets for input and output for each batch. + // Step 2c: Compute offsets for input and output for each batch. int64_t input_offset = descriptor.nside * descriptor.nside * 12; int64_t output_offset = (4 * descriptor.nside - 1) * (2 * descriptor.harmonic_band_limit); - // Step 2c: Fork CUDA streams for parallel processing of batches. + // Step 2d: Fork CUDA streams for parallel processing of batches. CudaStreamHandler handler; - handler.Fork(stream, batch_count); - auto stream_iter = handler.getIterator(); + std::vector sub_streams = handler.Fork(device_ordinal, stream, batch_count); - // Step 2d: Iterate over each batch. - for (int i = 0; i < batch_count && stream_iter.hasNext(); ++i) { - cudaStream_t sub_stream = stream_iter.next(); - // Step 2e: Get or create an s2fftExec instance from the PlanCache. + // Step 2e: Iterate over each batch. + for (int i = 0; i < batch_count; ++i) { + cudaStream_t sub_stream = sub_streams[i]; + // Step 2f: Get or create an s2fftExec instance from the PlanCache. auto executor = std::make_shared>(); - PlanCache::GetInstance().GetS2FFTExec(descriptor, executor); + PlanCache::GetInstance().GetS2FFTExec(descriptor, executor, device_ordinal); // Step 2f: Calculate device pointers for the current batch's data, output, and workspace. fft_complex_type* data_c = @@ -109,9 +140,9 @@ ffi::Error healpix_forward(cudaStream_t stream, ffi::Buffer input, ffi::Resul fft_complex_type* workspace_c = reinterpret_cast(workspace->typed_data() + i * executor->m_work_size); - // Step 2g: Launch the forward transform on this sub-stream. + // Step 2i: Launch the forward transform on this sub-stream. executor->Forward(descriptor, sub_stream, data_c, workspace_c); - // Step 2h: Launch spectral extension kernel with shift and normalization. + // Step 2j: Launch spectral extension kernel with shift and normalization. int kernel_norm = (descriptor.norm == s2fftKernels::fft_norm::FORWARD) ? 0 : (descriptor.norm == s2fftKernels::fft_norm::ORTHO) ? 1 : 2; @@ -119,22 +150,22 @@ ffi::Error healpix_forward(cudaStream_t stream, ffi::Buffer input, ffi::Resul descriptor.harmonic_band_limit, descriptor.shift, kernel_norm, sub_stream); } - // Step 2i: Join all forked streams back to the main stream. - handler.join(stream); + // Step 2k: Join all forked streams back to the main stream. + handler.join(stream, sub_streams); return ffi::Error::Success(); } else { - // Step 2j: Non-batched case. - // Step 2k: Get device pointers for data, output, and workspace. + // Step 2l: Non-batched case. + // Step 2m: Get device pointers for data, output, and workspace. fft_complex_type* data_c = reinterpret_cast(input_alias->typed_data()); fft_complex_type* out_c = reinterpret_cast(output->typed_data()); fft_complex_type* workspace_c = reinterpret_cast(workspace->typed_data()); - // Step 2l: Get or create an s2fftExec instance from the PlanCache. + // Step 2n: Get or create an s2fftExec instance from the PlanCache. auto executor = std::make_shared>(); - PlanCache::GetInstance().GetS2FFTExec(descriptor, executor); - // Step 2m: Launch the forward transform. + PlanCache::GetInstance().GetS2FFTExec(descriptor, executor, device_ordinal); + // Step 2o: Launch the forward transform. executor->Forward(descriptor, stream, data_c, workspace_c); - // Step 2n: Launch spectral extension kernel with shift and normalization. + // Step 2p: Launch spectral extension kernel with shift and normalization. int kernel_norm = (descriptor.norm == s2fftKernels::fft_norm::FORWARD) ? 0 : (descriptor.norm == s2fftKernels::fft_norm::ORTHO) ? 1 : 2; @@ -155,7 +186,7 @@ ffi::Error healpix_forward(cudaStream_t stream, ffi::Buffer input, ffi::Resul * * @tparam T The XLA data type. * @param stream CUDA stream to use. - * @param scratch ScratchAllocator for temporary device memory. + * @param device_ordinal CUDA device the buffers and plans live on. * @param input Input buffer containing FTM data. * @param output Output buffer to store HEALPix pixel-space data. * @param workspace Output buffer for temporary workspace memory. @@ -163,7 +194,7 @@ ffi::Error healpix_forward(cudaStream_t stream, ffi::Buffer input, ffi::Resul * @return ffi::Error indicating success or failure. */ template -ffi::Error healpix_backward(cudaStream_t stream, ffi::Buffer input, +ffi::Error healpix_backward(cudaStream_t stream, int32_t device_ordinal, ffi::Buffer input, ffi::Result> input_alias, ffi::Result> output, ffi::Result> workspace, s2fftDescriptor descriptor) { // Step 1: Determine the complex type based on the XLA data type. @@ -171,31 +202,35 @@ ffi::Error healpix_backward(cudaStream_t stream, ffi::Buffer input, const auto& dim_in = input.dimensions(); const auto& dim_out = output->dimensions(); - // Step 1a: Get shift strategy from descriptor. - bool is_batched = (dim_in.size() == 3); + // Step 1a: Make the input data visible in the aliased result buffer (no-op when they alias). + if (ffi::Error err = copy_operand_to_alias(stream, input, input_alias); !err.success()) { + return err; + } // Step 2: Handle batched and non-batched cases separately. + bool is_batched = (dim_in.size() == 3); + + // Step 2a: Handle batched and non-batched cases separately. if (is_batched) { - // Step 2a: Batched case. + // Step 2b: Batched case. // Assertions to ensure correct input/output dimensions for batched operations. assert(dim_out.size() == 2); assert(dim_in[0] == dim_out[0]); int batch_count = dim_in[0]; - // Step 2b: Compute offsets for input, output, and callback parameters for each batch. + // Step 2c: Compute offsets for input, output, and callback parameters for each batch. int64_t input_offset = (4 * descriptor.nside - 1) * (2 * descriptor.harmonic_band_limit); int64_t output_offset = descriptor.nside * descriptor.nside * 12; - // Step 2c: Fork CUDA streams for parallel processing of batches. + // Step 2d: Fork CUDA streams for parallel processing of batches. CudaStreamHandler handler; - handler.Fork(stream, batch_count); - auto stream_iter = handler.getIterator(); + std::vector sub_streams = handler.Fork(device_ordinal, stream, batch_count); - // Step 2d: Iterate over each batch. - for (int i = 0; i < batch_count && stream_iter.hasNext(); ++i) { - cudaStream_t sub_stream = stream_iter.next(); - // Step 2e: Get or create an s2fftExec instance from the PlanCache. + // Step 2e: Iterate over each batch. + for (int i = 0; i < batch_count; ++i) { + cudaStream_t sub_stream = sub_streams[i]; + // Step 2f: Get or create an s2fftExec instance from the PlanCache. auto executor = std::make_shared>(); - PlanCache::GetInstance().GetS2FFTExec(descriptor, executor); + PlanCache::GetInstance().GetS2FFTExec(descriptor, executor, device_ordinal); // Step 2f: Calculate device pointers for the current batch's data, output, and workspace. fft_complex_type* data_c = @@ -217,7 +252,7 @@ ffi::Error healpix_backward(cudaStream_t stream, ffi::Buffer input, executor->Backward(descriptor, sub_stream, out_c, workspace_c); } // Step 2i: Join all forked streams back to the main stream. - handler.join(stream); + handler.join(stream, sub_streams); return ffi::Error::Success(); } else { // Step 2j: Non-batched case. @@ -234,7 +269,7 @@ ffi::Error healpix_backward(cudaStream_t stream, ffi::Buffer input, // Step 2l: Get or create an s2fftExec instance from the PlanCache. auto executor = std::make_shared>(); - PlanCache::GetInstance().GetS2FFTExec(descriptor, executor); + PlanCache::GetInstance().GetS2FFTExec(descriptor, executor, device_ordinal); // Step 2m: Launch spectral folding kernel. s2fftKernels::launch_spectral_folding(data_c, out_c, descriptor.nside, descriptor.harmonic_band_limit, descriptor.shift, kernel_norm, stream); @@ -258,11 +293,13 @@ ffi::Error healpix_backward(cudaStream_t stream, ffi::Buffer input, * @param normalize Flag for normalization. * @param adjoint Flag indicating if an adjoint operation is desired. * @param must_exist If true, throws an error if the plan does not exist in the cache. + * @param device_ordinal The CUDA device the plans must belong to (see plan_cache.h). * @return s2fftDescriptor configured with the given parameters. */ template s2fftDescriptor build_descriptor(int64_t nside, int64_t harmonic_band_limit, bool reality, bool forward, - bool normalize, bool adjoint, bool must_exist, size_t& work_size) { + bool normalize, bool adjoint, bool must_exist, size_t& work_size, + int32_t device_ordinal) { using fft_complex_type = fft_complex_t; // Step 1: Determine FFT normalization type based on forward/normalize flags. s2fftKernels::fft_norm norm = s2fftKernels::fft_norm::NONE; @@ -281,10 +318,11 @@ s2fftDescriptor build_descriptor(int64_t nside, int64_t harmonic_band_limit, boo s2fftDescriptor descriptor(nside, harmonic_band_limit, reality, adjoint, forward, norm, shift, is_double_v); - // Step 4: Get or create an s2fftExec instance from the PlanCache. - // This call will also initialize the executor if it's newly created. + // Step 4: Get or create an s2fftExec instance from the device's PlanCache. + // This call will also initialize the executor if it's newly created; the caller holds the + // device guard, so newly created cuFFT plans bind to the right device. auto executor = std::make_shared>(); - HRESULT hr = PlanCache::GetInstance().GetS2FFTExec(descriptor, executor); + HRESULT hr = PlanCache::GetInstance().GetS2FFTExec(descriptor, executor, device_ordinal); // Step 5: Handle cases where the plan was expected to exist but didn't. if (hr == S_OK && must_exist) { // This is an error because S_OK means plan was created, but must_exist implies it should have been @@ -312,6 +350,7 @@ s2fftDescriptor build_descriptor(int64_t nside, int64_t harmonic_band_limit, boo * * @tparam T The XLA data type. * @param stream CUDA stream to use. + * @param device_ordinal CUDA device the buffers and plans belong to (XLA execution context). * @param nside HEALPix resolution parameter. * @param harmonic_band_limit Harmonic band limit L. * @param reality Flag indicating whether data is real-valued. @@ -324,20 +363,25 @@ s2fftDescriptor build_descriptor(int64_t nside, int64_t harmonic_band_limit, boo * @return ffi::Error indicating success or failure. */ template -ffi::Error healpix_fft_cuda(cudaStream_t stream, int64_t nside, int64_t harmonic_band_limit, bool reality, - bool forward, bool normalize, bool adjoint, ffi::Buffer input, - ffi::Result> input_alias, ffi::Result> output, - ffi::Result> workspace) { +ffi::Error healpix_fft_cuda(cudaStream_t stream, int32_t device_ordinal, int64_t nside, + int64_t harmonic_band_limit, bool reality, bool forward, bool normalize, + bool adjoint, ffi::Buffer input, ffi::Result> input_alias, + ffi::Result> output, ffi::Result> workspace) { + // Step 0: Select the device the buffers and plans belong to. The FFI context carries the + // device ordinal of the device executing this call; cuFFT plans created below bind to the + // device that is current at creation time (astro-informatics/s2fft#389). + CudaDeviceScope device_guard(device_ordinal); + // Step 1: Build the s2fftDescriptor based on the input parameters. size_t work_size = 0; // Variable to hold the workspace size s2fftDescriptor descriptor = build_descriptor(nside, harmonic_band_limit, reality, forward, normalize, - adjoint, true, work_size); + adjoint, false, work_size, device_ordinal); // Step 2: Dispatch to either forward or backward transform based on the 'forward' flag. if (forward) { - return healpix_forward(stream, input, input_alias, output, workspace, descriptor); + return healpix_forward(stream, device_ordinal, input, input_alias, output, workspace, descriptor); } else { - return healpix_backward(stream, input, input_alias, output, workspace, descriptor); + return healpix_backward(stream, device_ordinal, input, input_alias, output, workspace, descriptor); } } @@ -350,6 +394,7 @@ ffi::Error healpix_fft_cuda(cudaStream_t stream, int64_t nside, int64_t harmonic XLA_FFI_DEFINE_HANDLER_SYMBOL(healpix_fft_cuda_C64, healpix_fft_cuda, ffi::Ffi::Bind() .Ctx>() + .Ctx() .Attr("nside") .Attr("harmonic_band_limit") .Attr("reality") @@ -364,6 +409,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(healpix_fft_cuda_C64, healpix_fft_cuda, ffi::Ffi::Bind() .Ctx>() + .Ctx() .Attr("nside") .Attr("harmonic_band_limit") .Attr("reality") @@ -419,20 +465,28 @@ NB_MODULE(_s2fft, m) { m.def("registration", &s2fft::Registration); // Step 2: Declare and expose build_descriptor functions for C64 and C128 to Python. // These functions allow Python to query the required workspace size for a given descriptor. + // The warm-up executor is cached under the caller's current device so a runtime call on the + // same device reuses it; other devices create their own executor on first use. m.def("build_descriptor_C64", [](int64_t nside, int64_t harmonic_band_limit, bool reality, bool forward, bool normalize, bool adjoint) { - // Step 2a: Build the s2fftDescriptor. + // Step 2a: Build the s2fftDescriptor on the caller's current device. + int device = -1; + cudaGetDevice(&device); + CudaDeviceScope device_guard(device); size_t work_size = 0; // Variable to hold the workspace size s2fft::s2fftDescriptor desc = s2fft::build_descriptor( - nside, harmonic_band_limit, reality, forward, normalize, adjoint, false, work_size); + nside, harmonic_band_limit, reality, forward, normalize, adjoint, false, work_size, device); return work_size; }); m.def("build_descriptor_C128", [](int64_t nside, int64_t harmonic_band_limit, bool reality, bool forward, bool normalize, bool adjoint) { - // Step 2e: Build the s2fftDescriptor. + // Step 2e: Build the s2fftDescriptor on the caller's current device. + int device = -1; + cudaGetDevice(&device); + CudaDeviceScope device_guard(device); size_t work_size = 0; // Variable to hold the workspace size s2fft::s2fftDescriptor desc = s2fft::build_descriptor( - nside, harmonic_band_limit, reality, forward, normalize, adjoint, false, work_size); + nside, harmonic_band_limit, reality, forward, normalize, adjoint, false, work_size, device); return work_size; }); // Step 3: Expose a boolean attribute indicating if CUDA support is compiled in. diff --git a/lib/src/plan_cache.cc b/lib/src/plan_cache.cc index 1dd34cb5c..b5da32ea4 100644 --- a/lib/src/plan_cache.cc +++ b/lib/src/plan_cache.cc @@ -29,16 +29,19 @@ PlanCache::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 (see plan_cache.h). * @return HRESULT indicating success (S_OK if new, S_FALSE if from cache) or failure. */ template -HRESULT PlanCache::GetS2FFTExec(s2fftDescriptor &descriptor, std::shared_ptr> &executor) { +HRESULT PlanCache::GetS2FFTExec(s2fftDescriptor &descriptor, std::shared_ptr> &executor, + int32_t device_ordinal) { // Step 1: Check if the type is cufftComplex (single precision). if constexpr (std::is_same_v) { HRESULT hr(E_FAIL); - // Step 1a: Try to find the descriptor in the single-precision cache. - auto it = m_Descriptors32.find(descriptor); - if (it != m_Descriptors32.end()) { + // Step 1a: Try to find the descriptor in the device's single-precision cache. + auto &cache = m_Descriptors32[device_ordinal]; + auto it = cache.find(descriptor); + if (it != cache.end()) { // Step 1b: If found, retrieve the existing executor and set HR to S_FALSE (found in cache). executor = it->second; hr = S_FALSE; @@ -46,20 +49,22 @@ HRESULT PlanCache::GetS2FFTExec(s2fftDescriptor &descriptor, std::shared_ptrInitialize(descriptor); // Step 1e: If initialization is successful, store the new executor in the cache. if (SUCCEEDED(hr)) { - m_Descriptors32[descriptor] = executor; + cache[descriptor] = executor; } } // Step 1f: Return the HRESULT. return hr; } else { // Step 2: If the type is not cufftComplex, it must be cufftDoubleComplex (double precision). HRESULT hr(E_FAIL); - // Step 2a: Try to find the descriptor in the double-precision cache. - auto it = m_Descriptors64.find(descriptor); - if (it != m_Descriptors64.end()) { + // Step 2a: Try to find the descriptor in the device's double-precision cache. + auto &cache = m_Descriptors64[device_ordinal]; + auto it = cache.find(descriptor); + if (it != cache.end()) { // Step 2b: If found, retrieve the existing executor and set HR to S_FALSE (found in cache). executor = it->second; hr = S_FALSE; @@ -67,11 +72,12 @@ HRESULT PlanCache::GetS2FFTExec(s2fftDescriptor &descriptor, std::shared_ptrInitialize(descriptor); // Step 2e: If initialization is successful, store the new executor in the cache. if (SUCCEEDED(hr)) { - m_Descriptors64[descriptor] = executor; + cache[descriptor] = executor; } } // Step 2f: Return the HRESULT. @@ -110,9 +116,11 @@ PlanCache::~PlanCache() { // Explicitly instantiate the templates for the supported complex types. // This is necessary for the linker to find the concrete implementations of the templated function. template HRESULT PlanCache::GetS2FFTExec(s2fftDescriptor &descriptor, - std::shared_ptr> &executor); + std::shared_ptr> &executor, + int32_t device_ordinal); template HRESULT PlanCache::GetS2FFTExec( - s2fftDescriptor &descriptor, std::shared_ptr> &executor); + s2fftDescriptor &descriptor, std::shared_ptr> &executor, + int32_t device_ordinal); } // namespace s2fft \ No newline at end of file From a8dd5370f5aeb0f56960bb8c7c102ff53db83305 Mon Sep 17 00:00:00 2001 From: Wassim Kabalan Date: Sun, 13 Sep 2026 17:19:44 +0200 Subject: [PATCH 3/4] Update CMake to include SM90 for H100 GPUs --- CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 7848e1fd3..6e45bc7de 100644 --- a/CMakeLists.txt +++ b/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 From f2dfe5e09b5bb3f82b8ab23a4762c43730c0930a Mon Sep 17 00:00:00 2001 From: Wassim Kabalan Date: Sun, 13 Sep 2026 17:41:33 +0200 Subject: [PATCH 4/4] Update examples --- notebooks/JAX_CUDA_HEALPix.ipynb | 26 +- notebooks/MultiGPU_JAX_CUDA_Healpix.ipynb | 358 ++++++++++++++++++++++ 2 files changed, 371 insertions(+), 13 deletions(-) create mode 100644 notebooks/MultiGPU_JAX_CUDA_Healpix.ipynb diff --git a/notebooks/JAX_CUDA_HEALPix.ipynb b/notebooks/JAX_CUDA_HEALPix.ipynb index 397be6c77..78095b60e 100644 --- a/notebooks/JAX_CUDA_HEALPix.ipynb +++ b/notebooks/JAX_CUDA_HEALPix.ipynb @@ -54,7 +54,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "JAX version: 0.9.0.1\n", + "JAX version: 0.11.1\n", "JAX backend: gpu\n" ] } @@ -162,7 +162,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 5, "metadata": {}, "outputs": [ { @@ -206,21 +206,21 @@ "Forward Transform - First run (includes JIT compilation):\n", "\n", "CUDA:\n", - "CPU times: user 5.87 ms, sys: 886 μs, total: 6.76 ms\n", - "Wall time: 13.5 ms\n", + "CPU times: user 7.97 ms, sys: 0 ns, total: 7.97 ms\n", + "Wall time: 9.74 ms\n", "\n", "Pure JAX:\n", - "CPU times: user 2.98 s, sys: 116 ms, total: 3.1 s\n", - "Wall time: 2.45 s\n", + "CPU times: user 10.2 s, sys: 311 ms, total: 10.6 s\n", + "Wall time: 4.46 s\n", "\n", "============================================================\n", "Forward Transform - Execution time (after JIT):\n", "\n", "CUDA:\n", - "10.1 ms ± 153 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n", + "8.02 ms ± 28.1 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n", "\n", "Pure JAX:\n", - "10.4 ms ± 60.8 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" + "8.05 ms ± 113 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" ] } ], @@ -264,8 +264,8 @@ "output_type": "stream", "text": [ "Forward transform comparison:\n", - " Mean Squared Error: 1.28e-35\n", - " Max absolute difference: 2.86e-17\n", + " Mean Squared Error: 1.68e-35\n", + " Max absolute difference: 2.43e-17\n", " Results match: True\n" ] } @@ -428,9 +428,9 @@ ], "metadata": { "kernelspec": { - "display_name": "s2fft", + "display_name": "Particle Mesh", "language": "python", - "name": "python3" + "name": "pm" }, "language_info": { "codemirror_mode": { @@ -442,7 +442,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.4" + "version": "3.12.7" } }, "nbformat": 4, diff --git a/notebooks/MultiGPU_JAX_CUDA_Healpix.ipynb b/notebooks/MultiGPU_JAX_CUDA_Healpix.ipynb new file mode 100644 index 000000000..f81314891 --- /dev/null +++ b/notebooks/MultiGPU_JAX_CUDA_Healpix.ipynb @@ -0,0 +1,358 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "title-md", + "metadata": {}, + "source": [ + "# Multi-GPU HEALPix transforms with the s2fft CUDA backend\n", + "\n", + "Sharded `s2fft.forward` (map to harmonics) over a JAX device mesh, comparing the\n", + "`jax_cuda` custom-CUDA path against the pure-JAX path. Uses only `s2fft` — no `jax_healpy`.\n", + "\n", + "## Before running on the cluster\n", + "\n", + "This notebook requires the **ASKabalan fork build** of `s2fft` (per-device cuFFT plan\n", + "cache + stream pools, fixed FFI aliasing; see `HANDOFF.md`). The verification cell asserts\n", + "the fork version and runs a single-GPU `jax_cuda` sanity check. If the cluster venv holds\n", + "the upstream package, install the fork first (wheel/git — delivery of the build to the\n", + "cluster is handled outside this notebook).\n", + "\n", + "If cuFFT plan creation fails with code 5 on small GPUs, cap JAX's pool:\n", + "`XLA_PYCLIENT_MEM_FRACTION<=0.8` (export before starting the kernel)." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "setup", + "metadata": { + "tags": [ + "parameters" + ] + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "JAX is not using 64-bit precision. This will dramatically affect numerical precision at even moderate L.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Mesh shape : OrderedDict({'x': 4, 'y': 1})\n", + "Partition spec : P('x', 'y')\n" + ] + } + ], + "source": [ + "from functools import partial\n", + "\n", + "import jax\n", + "import jax.numpy as jnp\n", + "import jax.random as jr\n", + "import numpy as np\n", + "from jax.sharding import AxisType, NamedSharding\n", + "from jax.sharding import PartitionSpec as P\n", + "\n", + "import s2fft\n", + "from s2fft.utils import signal_generator\n", + "\n", + "jax.config.update(\"jax_enable_x64\", True)\n", + "\n", + "nside = 32\n", + "L = 2 * nside\n", + "key = jr.key(0)\n", + "\n", + "mesh = jax.make_mesh((jax.device_count(), 1), (\"x\", \"y\"), axis_types=(AxisType.Auto, AxisType.Auto))\n", + "sharding = NamedSharding(mesh, P(\"x\", \"y\"))\n", + "print(f\"Mesh shape : {mesh.shape}\")\n", + "print(f\"Partition spec : {sharding.spec}\")" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "verify", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "s2fft 1.4.1.dev21 | CUDA extension OK | jax 0.11.1 | devices: [CudaDevice(id=0), CudaDevice(id=1), CudaDevice(id=2), CudaDevice(id=3)]\n", + "single-GPU forward sanity: max|a_cuda-a_jax|/max|a_jax| = 1.55e-16\n" + ] + } + ], + "source": [ + "# Environment verification: this notebook requires the ASKabalan fork build of s2fft\n", + "# (per-device cuFFT plan cache + stream pools, fixed FFI aliasing). If this cell fails,\n", + "# the venv still holds the upstream package and every jax_cuda cell below is void.\n", + "import importlib.metadata as md\n", + "\n", + "import s2fft_lib._s2fft as _ext\n", + "\n", + "version = md.version(\"s2fft\")\n", + "assert version.startswith(\"1.4.1.dev\"), f\"upstream s2fft {version} installed - install the fork build first\"\n", + "assert _ext.COMPILED_WITH_CUDA, \"s2fft_lib was compiled without CUDA support\"\n", + "print(f\"s2fft {version} | CUDA extension OK | jax {jax.__version__} | devices: {jax.devices()}\")\n", + "\n", + "# Tiny single-GPU sanity check through the CUDA path (agrees with the jax path).\n", + "nside_s = 8\n", + "L_s = 2 * nside_s\n", + "rng = np.random.default_rng(0)\n", + "flm_s = signal_generator.generate_flm(rng=rng, L=L_s, reality=True)\n", + "f_s = s2fft.inverse(flm_s, L=L_s, nside=nside_s, sampling=\"healpix\", reality=True, method=\"jax\")\n", + "alm_c = s2fft.forward(f_s, L=L_s, nside=nside_s, sampling=\"healpix\", reality=True, method=\"jax_cuda\")\n", + "alm_j = s2fft.forward(f_s, L=L_s, nside=nside_s, sampling=\"healpix\", reality=True, method=\"jax\")\n", + "rel = float(jnp.max(jnp.abs(alm_c - alm_j)) / jnp.max(jnp.abs(alm_j)))\n", + "print(f\"single-GPU forward sanity: max|a_cuda-a_jax|/max|a_jax| = {rel:.2e}\")\n", + "assert rel < 1e-10, \"jax_cuda disagrees with jax on a single GPU\"" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "id": "data", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
                                                                                \n",
+       "                                     GPU 0                                      \n",
+       "                                                                                \n",
+       "                                                                                \n",
+       "                                     GPU 1                                      \n",
+       "                                                                                \n",
+       "                                                                                \n",
+       "                                     GPU 2                                      \n",
+       "                                                                                \n",
+       "                                                                                \n",
+       "                                     GPU 3                                      \n",
+       "                                                                                \n",
+       "
\n" + ], + "text/plain": [ + "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\n", + "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121mGPU 0\u001b[0m\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\n", + "\u001b[38;2;255;255;255;48;2;57;59;121m \u001b[0m\n", + "\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\n", + "\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214mGPU 1\u001b[0m\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\n", + "\u001b[38;2;255;255;255;48;2;222;158;214m \u001b[0m\n", + "\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\n", + "\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74mGPU 2\u001b[0m\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\n", + "\u001b[38;2;255;255;255;48;2;173;73;74m \u001b[0m\n", + "\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", + "\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107mGPU 3\u001b[0m\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n", + "\u001b[38;2;0;0;0;48;2;181;207;107m \u001b[0m\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "npix = 12 * nside**2\n", + "\n", + "# One random real sky per device, generated through the jax path.\n", + "flm_per_device = [\n", + " signal_generator.generate_flm(rng=np.random.default_rng(d), L=L, reality=True)\n", + " for d in range(jax.device_count())\n", + "]\n", + "f_per_device = [\n", + " s2fft.inverse(flm, L=L, nside=nside, sampling=\"healpix\", reality=True, method=\"jax\")\n", + " for flm in flm_per_device\n", + "]\n", + "hp_map = jnp.asarray(np.stack(f_per_device))\n", + "hp_map = jax.lax.with_sharding_constraint(hp_map, sharding)\n", + "\n", + "\n", + "def sharded_forward(f, method):\n", + " def inner_fn(f):\n", + " return s2fft.forward(\n", + " f.squeeze(),\n", + " L=L,\n", + " nside=nside,\n", + " sampling=\"healpix\",\n", + " reality=False,\n", + " method=method,\n", + " )\n", + " return jax.shard_map(inner_fn, mesh=mesh, in_specs=P(\"x\", None), out_specs=P(\"x\", None))(f)\n", + "\n", + "jax.debug.visualize_array_sharding(hp_map)" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "clear", + "metadata": {}, + "outputs": [], + "source": [ + "jax.clear_caches()" + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "id": "time-cuda", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 2.73 s, sys: 97.7 ms, total: 2.83 s\n", + "Wall time: 1.67 s\n" + ] + } + ], + "source": [ + "%time cuda_flm = sharded_forward(hp_map, method=\"jax_cuda\").block_until_ready()" + ] + }, + { + "cell_type": "code", + "execution_count": 26, + "id": "time-jax", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 10.5 s, sys: 358 ms, total: 10.8 s\n", + "Wall time: 4.98 s\n" + ] + } + ], + "source": [ + "%time jax_flm = sharded_forward(hp_map, method=\"jax\").block_until_ready()" + ] + }, + { + "cell_type": "code", + "execution_count": 27, + "id": "determinism", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "jax_cuda determinism across runs: max|run1 - run2| = 0.000e+00\n" + ] + } + ], + "source": [ + "# Determinism: two independent jax_cuda sharded runs must agree bit-for-bit.\n", + "cuda_flm_1 = sharded_forward(hp_map, method=\"jax_cuda\").block_until_ready()\n", + "cuda_flm_2 = sharded_forward(hp_map, method=\"jax_cuda\").block_until_ready()\n", + "max_diff = float(jnp.max(jnp.abs(cuda_flm_1 - cuda_flm_2)))\n", + "print(f\"jax_cuda determinism across runs: max|run1 - run2| = {max_diff:.3e}\")\n", + "assert max_diff == 0.0, \"jax_cuda sharded forward is not deterministic\"" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "id": "agreement", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "sharded jax_cuda vs jax: max|a_cuda-a_jax|/max|a_jax| = 3.34e-16\n" + ] + } + ], + "source": [ + "# Agreement: sharded jax_cuda vs sharded jax on identical data (float64).\n", + "rel = float(jnp.max(jnp.abs(cuda_flm_1 - jax_flm)) / jnp.max(jnp.abs(jax_flm)))\n", + "print(f\"sharded jax_cuda vs jax: max|a_cuda-a_jax|/max|a_jax| = {rel:.2e}\")\n", + "assert rel < 1e-10, \"sharded jax_cuda disagrees with sharded jax\"" + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "id": "timeit-cuda", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1.23 s ± 19.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" + ] + } + ], + "source": [ + "%timeit cuda_flm = sharded_forward(hp_map, method=\"jax_cuda\").block_until_ready()" + ] + }, + { + "cell_type": "code", + "execution_count": 30, + "id": "timeit-jax", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "4.49 s ± 22.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" + ] + } + ], + "source": [ + "%timeit jax_flm = sharded_forward(hp_map, method=\"jax\").block_until_ready()" + ] + }, + { + "cell_type": "code", + "execution_count": 31, + "id": "final", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Finished all good\n" + ] + } + ], + "source": [ + "print(\"Finished all good\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Particle Mesh", + "language": "python", + "name": "pm" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.7" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}