From 5d1700493a2d199e4f3aac7c5a8df4925c05039f Mon Sep 17 00:00:00 2001 From: Lucas Newman Date: Sat, 14 Mar 2026 05:02:19 -0700 Subject: [PATCH] [CUDA] Add FFT support (#3243) --- .github/actions/test-linux/action.yml | 2 +- mlx/backend/cuda/CMakeLists.txt | 4 + mlx/backend/cuda/fft.cu | 443 ++++++++++++++++++++++++++ mlx/backend/cuda/primitives.cpp | 1 - mlx/backend/gpu/copy.cpp | 34 +- mlx/backend/gpu/copy.h | 2 + python/tests/cuda_skip.py | 10 - python/tests/test_fft.py | 11 + tests/fft_tests.cpp | 15 + 9 files changed, 498 insertions(+), 24 deletions(-) create mode 100644 mlx/backend/cuda/fft.cu diff --git a/.github/actions/test-linux/action.yml b/.github/actions/test-linux/action.yml index 3258670e..b2d3e17d 100644 --- a/.github/actions/test-linux/action.yml +++ b/.github/actions/test-linux/action.yml @@ -65,5 +65,5 @@ runs: DEVICE: gpu run: | echo "::group::CPP tests - GPU" - ./build/tests/tests -sfe="*fft_tests.cpp,*linalg_tests.cpp" + ./build/tests/tests -sfe="*linalg_tests.cpp" echo "::endgroup::" diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 5cac9e59..08d1ba1d 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -26,6 +26,7 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cu ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp ${CMAKE_CURRENT_SOURCE_DIR}/event.cu + ${CMAKE_CURRENT_SOURCE_DIR}/fft.cu ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gemv.cu ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm.cpp @@ -247,6 +248,9 @@ target_include_directories(mlx PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) # Use cublasLt. target_link_libraries(mlx PRIVATE CUDA::cublasLt) +# Use cuFFT. +target_link_libraries(mlx PRIVATE CUDA::cufft) + # Use NVRTC and driver APIs. target_link_libraries(mlx PRIVATE CUDA::nvrtc CUDA::cuda_driver) diff --git a/mlx/backend/cuda/fft.cu b/mlx/backend/cuda/fft.cu new file mode 100644 index 00000000..a657bb2d --- /dev/null +++ b/mlx/backend/cuda/fft.cu @@ -0,0 +1,443 @@ +// Copyright © 2025 Apple Inc. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "mlx/backend/common/utils.h" +#include "mlx/backend/cuda/allocator.h" +#include "mlx/backend/cuda/device.h" +#include "mlx/backend/cuda/device/complex.cuh" +#include "mlx/backend/cuda/lru_cache.h" +#include "mlx/backend/cuda/utils.h" +#include "mlx/backend/gpu/copy.h" +#include "mlx/primitives.h" + +namespace mlx::core { + +namespace cu { + +namespace cg = cooperative_groups; + +template +__global__ void scale_fft_output(T* out, T scale, size_t size) { + auto index = cg::this_grid().thread_rank(); + if (index < size) { + out[index] *= scale; + } +} + +} // namespace cu + +namespace { + +void check_cufft_error(const char* name, cufftResult err) { + if (err != CUFFT_SUCCESS) { + throw std::runtime_error( + std::string(name) + + " failed with code: " + std::to_string(static_cast(err)) + "."); + } +} + +#define CHECK_CUFFT_ERROR(cmd) check_cufft_error(#cmd, (cmd)) + +enum class FFTTransformType : uint8_t { + C2C = 0, + R2C = 1, + C2R = 2, +}; + +struct FFTPlanKey { + int device_id; + FFTTransformType transform_type; + int64_t n; + int64_t batch; +}; + +struct CuFFTPlan { + explicit CuFFTPlan(int device_id, cufftHandle handle, size_t workspace_size) + : device_id(device_id), handle(handle), workspace_size(workspace_size) {} + + ~CuFFTPlan() { + if (handle != 0) { + try { + cu::device(device_id).make_current(); + cufftDestroy(handle); + } catch (...) { + } + } + } + + int device_id; + cufftHandle handle; + size_t workspace_size; +}; + +struct OrderedArray { + array arr; + std::vector order; +}; + +auto& fft_plan_cache() { + static LRUBytesKeyCache> cache( + "MLX_CUDA_FFT_CACHE_SIZE", + /* default_capacity */ 128); + return cache; +} + +FFTPlanKey make_plan_key( + int device_id, + FFTTransformType transform_type, + int64_t n, + int64_t batch) { + FFTPlanKey key{}; + key.device_id = device_id; + key.transform_type = transform_type; + key.n = n; + key.batch = batch; + return key; +} + +cudaDataType_t input_type(FFTTransformType transform_type) { + switch (transform_type) { + case FFTTransformType::C2C: + case FFTTransformType::C2R: + return CUDA_C_32F; + case FFTTransformType::R2C: + return CUDA_R_32F; + } + throw std::runtime_error("[FFT] Unsupported cuFFT input transform type."); +} + +cudaDataType_t output_type(FFTTransformType transform_type) { + switch (transform_type) { + case FFTTransformType::C2C: + case FFTTransformType::R2C: + return CUDA_C_32F; + case FFTTransformType::C2R: + return CUDA_R_32F; + } + throw std::runtime_error("[FFT] Unsupported cuFFT output transform type."); +} + +cudaDataType_t execution_type(FFTTransformType transform_type) { + switch (transform_type) { + case FFTTransformType::C2C: + return CUDA_C_32F; + case FFTTransformType::R2C: + return CUDA_R_32F; + case FFTTransformType::C2R: + return CUDA_C_32F; + } + throw std::runtime_error("[FFT] Unsupported cuFFT execution transform type."); +} + +int64_t input_embed(FFTTransformType transform_type, int64_t n) { + return transform_type == FFTTransformType::C2R ? (n / 2 + 1) : n; +} + +int64_t output_embed(FFTTransformType transform_type, int64_t n) { + return transform_type == FFTTransformType::R2C ? (n / 2 + 1) : n; +} + +int exec_direction(FFTTransformType transform_type, bool inverse) { + switch (transform_type) { + case FFTTransformType::C2C: + return inverse ? CUFFT_INVERSE : CUFFT_FORWARD; + case FFTTransformType::R2C: + return CUFFT_FORWARD; + case FFTTransformType::C2R: + return CUFFT_INVERSE; + } + throw std::runtime_error("[FFT] Unsupported cuFFT execution direction."); +} + +std::shared_ptr get_fft_plan( + cu::CommandEncoder& encoder, + FFTTransformType transform_type, + int64_t n, + int64_t batch) { + auto key = BytesKey{}; + key.pod = + make_plan_key(encoder.device().cuda_device(), transform_type, n, batch); + + auto& cache = fft_plan_cache(); + if (auto entry = cache.find(key); entry != cache.end()) { + return entry->second; + } + + encoder.device().make_current(); + + cufftHandle handle = 0; + size_t workspace_size = 0; + try { + CHECK_CUFFT_ERROR(cufftCreate(&handle)); + CHECK_CUFFT_ERROR(cufftSetAutoAllocation(handle, 0)); + CHECK_CUFFT_ERROR(cufftSetStream(handle, encoder.stream())); + + long long plan_n[1] = {n}; + long long inembed[1] = {input_embed(transform_type, n)}; + long long onembed[1] = {output_embed(transform_type, n)}; + CHECK_CUFFT_ERROR(cufftXtMakePlanMany( + handle, + /* rank= */ 1, + plan_n, + inembed, + /* istride= */ 1, + /* idist= */ input_embed(transform_type, n), + input_type(transform_type), + onembed, + /* ostride= */ 1, + /* odist= */ output_embed(transform_type, n), + output_type(transform_type), + batch, + &workspace_size, + execution_type(transform_type))); + } catch (...) { + if (handle != 0) { + encoder.device().make_current(); + cufftDestroy(handle); + } + throw; + } + + auto plan = std::make_shared( + encoder.device().cuda_device(), handle, workspace_size); + return cache.emplace(key, plan).first->second; +} + +std::vector make_identity_order(int ndim) { + std::vector order(ndim); + std::iota(order.begin(), order.end(), 0); + return order; +} + +std::vector move_axis_to_back_permutation(int ndim, int axis_pos) { + std::vector perm; + perm.reserve(ndim); + for (int i = 0; i < ndim; ++i) { + if (i != axis_pos) { + perm.push_back(i); + } + } + perm.push_back(axis_pos); + return perm; +} + +std::vector apply_permutation( + const std::vector& values, + const std::vector& perm) { + std::vector out(perm.size()); + for (int i = 0; i < perm.size(); ++i) { + out[i] = values[perm[i]]; + } + return out; +} + +int find_axis_position(const std::vector& order, int axis) { + auto it = std::find(order.begin(), order.end(), axis); + if (it == order.end()) { + throw std::runtime_error("[FFT] Internal axis tracking mismatch."); + } + return static_cast(it - order.begin()); +} + +OrderedArray prepare_input( + const OrderedArray& current, + int axis, + bool allow_direct, + cu::CommandEncoder& encoder, + Stream s) { + int axis_pos = find_axis_position(current.order, axis); + bool axis_last = axis_pos == static_cast(current.order.size()) - 1; + bool direct = allow_direct && axis_last && current.arr.flags().row_contiguous; + + if (direct) { + return current; + } + + array view = current.arr; + std::vector order = current.order; + if (!axis_last) { + auto perm = move_axis_to_back_permutation(current.arr.ndim(), axis_pos); + view = transpose_in_eval(current.arr, perm); + order = apply_permutation(current.order, perm); + } + + array packed = contiguous_copy_gpu(view, s); + encoder.add_temporary(packed); + return {std::move(packed), std::move(order)}; +} + +void execute_fft( + const array& in, + array& out, + FFTTransformType transform_type, + bool inverse, + cu::CommandEncoder& encoder) { + if (!in.flags().row_contiguous || in.strides(-1) != 1) { + throw std::runtime_error("[FFT] Expected packed row-contiguous FFT input."); + } + + int64_t n = + transform_type == FFTTransformType::C2R ? out.shape(-1) : in.shape(-1); + int64_t batch = in.shape().empty() ? 1 : in.size() / in.shape(-1); + auto plan = get_fft_plan(encoder, transform_type, n, batch); + + encoder.set_input_array(in); + out.set_data(cu::malloc_async(out.nbytes(), encoder)); + encoder.set_output_array(out); + encoder.add_completed_handler([plan]() {}); + + encoder.device().make_current(); + CHECK_CUFFT_ERROR(cufftSetStream(plan->handle, encoder.stream())); + auto* workspace = allocate_workspace(encoder, plan->workspace_size); + CHECK_CUFFT_ERROR(cufftSetWorkArea(plan->handle, workspace)); + + auto capture = encoder.capture_context(); + CHECK_CUFFT_ERROR(cufftXtExec( + plan->handle, + gpu_ptr(in), + gpu_ptr(out), + exec_direction(transform_type, inverse))); +} + +void restore_output_layout(const OrderedArray& current, array& out) { + Strides out_strides(out.ndim()); + for (int i = 0; i < current.order.size(); ++i) { + out_strides[current.order[i]] = current.arr.strides(i); + } + + auto [data_size, row_contiguous, col_contiguous] = + check_contiguity(out.shape(), out_strides); + bool contiguous = + current.arr.flags().contiguous && data_size == current.arr.data_size(); + + out.copy_shared_buffer( + current.arr, + out_strides, + {contiguous, row_contiguous, col_contiguous}, + current.arr.data_size()); +} + +void apply_inverse_scale( + array& arr, + const std::vector& axes, + const array& out, + cu::CommandEncoder& encoder) { + if (axes.empty()) { + return; + } + + double scale = 1.0; + for (auto axis : axes) { + scale /= out.shape(axis); + } + + size_t size = arr.data_size(); + dim3 block_dims(256); + dim3 grid_dims((size + block_dims.x - 1) / block_dims.x); + + encoder.set_input_array(arr); + encoder.set_output_array(arr); + + if (arr.dtype() == float32) { + float scale_f = static_cast(scale); + encoder.add_kernel_node( + cu::scale_fft_output, + grid_dims, + block_dims, + gpu_ptr(arr), + scale_f, + size); + } else if (arr.dtype() == complex64) { + cu::complex64_t scale_f(static_cast(scale), 0.0f); + encoder.add_kernel_node( + cu::scale_fft_output, + grid_dims, + block_dims, + gpu_ptr(arr), + scale_f, + size); + } else { + throw std::runtime_error("[FFT] Unsupported dtype for inverse scaling."); + } +} + +} // namespace + +void FFT::eval_gpu(const std::vector& inputs, array& out) { + nvtx3::scoped_range r("FFT::eval_gpu"); + auto& s = stream(); + auto& encoder = cu::get_command_encoder(s); + auto& in = inputs[0]; + + if (out.size() == 0) { + return; + } + + auto order = make_identity_order(in.ndim()); + OrderedArray current{in, std::move(order)}; + + std::vector axis_sequence; + axis_sequence.reserve(axes_.size()); + if (inverse_) { + for (auto axis : axes_) { + axis_sequence.push_back(static_cast(axis)); + } + } else { + for (int i = static_cast(axes_.size()) - 1; i >= 0; --i) { + axis_sequence.push_back(static_cast(axes_[i])); + } + } + + int real_axis = axes_.empty() ? -1 : static_cast(axes_.back()); + + for (int i = 0; i < axis_sequence.size(); ++i) { + int axis = axis_sequence[i]; + bool step_real = real_ && axis == real_axis; + auto transform_type = step_real + ? (inverse_ ? FFTTransformType::C2R : FFTTransformType::R2C) + : FFTTransformType::C2C; + + // cuFFT may overwrite the input buffer for C2R, so only use the direct + // input when the transform is out-of-place from the library's perspective + // or when the original input may be donated to the output. + auto prepared = prepare_input( + current, + axis, + /* allow_direct= */ transform_type != FFTTransformType::C2R || + is_donatable(in, out), + encoder, + s); + + Shape step_shape = prepared.arr.shape(); + if (step_real) { + step_shape.back() = out.shape(axis); + } + + Dtype step_dtype = + transform_type == FFTTransformType::C2R ? float32 : complex64; + array step_out(std::move(step_shape), step_dtype, nullptr, {}); + execute_fft(prepared.arr, step_out, transform_type, inverse_, encoder); + encoder.add_temporary(step_out); + + current = {std::move(step_out), std::move(prepared.order)}; + } + + if (inverse_) { + apply_inverse_scale(current.arr, axes_, out, encoder); + } + + restore_output_layout(current, out); +} + +} // namespace mlx::core diff --git a/mlx/backend/cuda/primitives.cpp b/mlx/backend/cuda/primitives.cpp index 1321ec38..0bf45214 100644 --- a/mlx/backend/cuda/primitives.cpp +++ b/mlx/backend/cuda/primitives.cpp @@ -25,7 +25,6 @@ namespace mlx::core { } NO_GPU(BlockMaskedMM) -NO_GPU(FFT) NO_GPU(GatherQMM) NO_GPU_MULTI(LUF) NO_GPU_MULTI(QRF) diff --git a/mlx/backend/gpu/copy.cpp b/mlx/backend/gpu/copy.cpp index 403d07d9..699dadd1 100644 --- a/mlx/backend/gpu/copy.cpp +++ b/mlx/backend/gpu/copy.cpp @@ -4,6 +4,7 @@ #include "mlx/primitives.h" #include +#include namespace mlx::core { @@ -59,19 +60,13 @@ array reshape_in_eval(const array& x, Shape shape, Stream s) { return out; } -array swapaxes_in_eval(const array& x, int axis1, int axis2) { - int ndim = x.ndim(); - if (axis1 < 0) { - axis1 += ndim; +array transpose_in_eval(const array& x, const std::vector& axes) { + Shape shape(axes.size()); + Strides strides(axes.size()); + for (int i = 0; i < axes.size(); ++i) { + shape[i] = x.shape(axes[i]); + strides[i] = x.strides(axes[i]); } - if (axis2 < 0) { - axis2 += ndim; - } - - auto shape = x.shape(); - std::swap(shape[axis1], shape[axis2]); - auto strides = x.strides(); - std::swap(strides[axis1], strides[axis2]); auto [data_size, row_contiguous, col_contiguous] = check_contiguity(shape, strides); @@ -86,4 +81,19 @@ array swapaxes_in_eval(const array& x, int axis1, int axis2) { return out; } +array swapaxes_in_eval(const array& x, int axis1, int axis2) { + int ndim = x.ndim(); + if (axis1 < 0) { + axis1 += ndim; + } + if (axis2 < 0) { + axis2 += ndim; + } + + std::vector axes(ndim); + std::iota(axes.begin(), axes.end(), 0); + std::swap(axes[axis1], axes[axis2]); + return transpose_in_eval(x, axes); +} + } // namespace mlx::core diff --git a/mlx/backend/gpu/copy.h b/mlx/backend/gpu/copy.h index 6e6bc797..8ebcd846 100644 --- a/mlx/backend/gpu/copy.h +++ b/mlx/backend/gpu/copy.h @@ -6,6 +6,7 @@ #include "mlx/stream.h" #include +#include namespace mlx::core { @@ -52,6 +53,7 @@ void reshape_gpu(const array& in, array& out, Stream s); // Like the normal ops but safe to call in eval_gpu. array flatten_in_eval(const array& x, int start_axis, int end_axis, Stream s); array reshape_in_eval(const array& x, Shape shape, Stream s); +array transpose_in_eval(const array& x, const std::vector& axes); array swapaxes_in_eval(const array& x, int axis1, int axis2); } // namespace mlx::core diff --git a/python/tests/cuda_skip.py b/python/tests/cuda_skip.py index 888be52b..6de59455 100644 --- a/python/tests/cuda_skip.py +++ b/python/tests/cuda_skip.py @@ -6,16 +6,6 @@ cuda_skip = { "TestBlas.test_gather_matmul", "TestBlas.test_gather_matmul_grad", "TestBlas.test_gather_mm_sorted_vjp", - # FFTs NYI - "TestFFT.test_fft", - "TestFFT.test_fft_big_powers_of_two", - "TestFFT.test_fft_contiguity", - "TestFFT.test_fft_exhaustive", - "TestFFT.test_fft_grads", - "TestFFT.test_fft_into_ifft", - "TestFFT.test_fft_large_numbers", - "TestFFT.test_fft_shared_mem", - "TestFFT.test_fftn", # Lapack ops NYI "TestLinalg.test_cholesky", "TestLinalg.test_cholesky_inv", diff --git a/python/tests/test_fft.py b/python/tests/test_fft.py index 07ab6267..9921c7d9 100644 --- a/python/tests/test_fft.py +++ b/python/tests/test_fft.py @@ -91,6 +91,17 @@ class TestFFT(mlx_tests.MLXTestCase): np_op = getattr(np.fft, op) self.check_mx_np(mx_op, np_op, x, axes=ax, s=s) + # Explicitly exercise transposed layouts and axes that are not + # physically last in memory order. + xt = np.transpose(a, (1, 2, 0)) + self.check_mx_np(mx.fft.fftn, np.fft.fftn, xt, axes=(2, 0)) + self.check_mx_np(mx.fft.ifftn, np.fft.ifftn, xt, axes=(2, 0)) + + rt = np.transpose(r, (1, 2, 0)) + self.check_mx_np(mx.fft.rfftn, np.fft.rfftn, rt, axes=(2, 0)) + irfft_in = np.ascontiguousarray(np.fft.rfftn(rt, axes=(2, 0))) + self.check_mx_np(mx.fft.irfftn, np.fft.irfftn, irfft_in, axes=(2, 0)) + def _run_ffts(self, shape, atol=1e-4, rtol=1e-4): np.random.seed(9) diff --git a/tests/fft_tests.cpp b/tests/fft_tests.cpp index b9e2d1bc..39ed80bd 100644 --- a/tests/fft_tests.cpp +++ b/tests/fft_tests.cpp @@ -193,6 +193,21 @@ TEST_CASE("test fftn") { CHECK_EQ(y.shape(), Shape{5, 8}); CHECK_EQ(y.dtype(), float32); } + + // Test non-contiguous layouts and axes that are not physically last. + { + x = astype( + transpose(reshape(arange(24, float32), {2, 3, 4}), {1, 2, 0}), + complex64); + auto y = fft::fftn(x, {2, 0}); + CHECK_EQ(y.shape(), x.shape()); + CHECK(allclose(fft::ifftn(y, {2, 0}), x, 1e-5, 1e-5).item()); + + auto r = transpose(reshape(arange(60, float32), {3, 4, 5}), {1, 2, 0}); + auto yr = fft::rfftn(r, {2, 0}); + CHECK_EQ(yr.shape(), Shape{3, 5, 3}); + CHECK(allclose(fft::irfftn(yr, {2, 0}), r, 1e-5, 1e-5).item()); + } } TEST_CASE("test fft with provided shape") {