From b5dd6dcff415fcf33865452153507c9c3b03082c Mon Sep 17 00:00:00 2001 From: Zheming Jin Date: Mon, 27 Jul 2026 21:26:22 -0500 Subject: [PATCH] [DFT][cuFFT] Fix "Invalid strides" for shapes with a unit dimension When a DFT descriptor has a dimension of extent 1 (e.g. {3,1}, {2,3,1}), the default strides of that dimension tie with its neighbour's stride. The cuFFT commit used std::min_element to locate the smallest-stride (innermost) dimension, which returns the first tied element and thus mis-selects an outer dimension as innermost. This triggered a spurious dimension swap and made the check_stride_validity() pre-check reject a perfectly valid transform with "Invalid strides", even though the underlying cufftPlanMany call succeeds. Search for the smallest stride in reverse so ties resolve to the innermost (highest-index) dimension, avoiding the incorrect swap. Also capture the selected index before erase() to avoid using an invalidated iterator. Fixes #631. --- src/dft/backends/cufft/commit.cpp | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/src/dft/backends/cufft/commit.cpp b/src/dft/backends/cufft/commit.cpp index 91cd17971..923122f08 100644 --- a/src/dft/backends/cufft/commit.cpp +++ b/src/dft/backends/cufft/commit.cpp @@ -25,6 +25,7 @@ #include #include +#include #include #include "oneapi/math/exceptions.hpp" @@ -160,9 +161,20 @@ class cufft_commit final : public dft::detail::commit_impl { offset_bwd_in = stride_vecs.offset_bwd_in; offset_bwd_out = stride_vecs.offset_bwd_out; - // cufft ignores the first value in inembed and onembed, so there is no harm in putting offset there - auto a_min = std::min_element(stride_vecs.vec_a.begin() + 1, stride_vecs.vec_a.end()); - auto b_min = std::min_element(stride_vecs.vec_b.begin() + 1, stride_vecs.vec_b.end()); + // cufft ignores the first value in inembed and onembed, so there is no harm in putting offset there. + // Find the dimension with the smallest stride (used as the cuFFT istride/ostride). When several + // dimensions share the smallest stride - which happens when a dimension has extent 1 and hence the + // same stride as its neighbour - prefer the innermost (last) dimension. Searching in reverse makes + // std::min_element resolve such ties to the highest index, avoiding a spurious dimension swap below + // that would otherwise reject valid transforms (e.g. shapes like {3,1}). See issue #631. + auto a_min = std::min_element(std::make_reverse_iterator(stride_vecs.vec_a.end()), + std::make_reverse_iterator(stride_vecs.vec_a.begin() + 1)) + .base() - + 1; + auto b_min = std::min_element(std::make_reverse_iterator(stride_vecs.vec_b.end()), + std::make_reverse_iterator(stride_vecs.vec_b.begin() + 1)) + .base() - + 1; if constexpr (dom == dft::domain::REAL) { if ((a_min != stride_vecs.vec_a.begin() + rank) || (b_min != stride_vecs.vec_b.begin() + rank)) { @@ -180,6 +192,8 @@ class cufft_commit final : public dft::detail::commit_impl { } const int a_stride = static_cast(*a_min); const int b_stride = static_cast(*b_min); + // Capture the position before erasing, as the iterator is invalidated by erase. + const auto a_min_idx = a_min - stride_vecs.vec_a.begin(); stride_vecs.vec_a.erase(a_min); stride_vecs.vec_b.erase(b_min); int fwd_istride = a_stride; @@ -188,9 +202,9 @@ class cufft_commit final : public dft::detail::commit_impl { stride_api_choice == dft::detail::stride_api::FB_STRIDES ? b_stride : a_stride; int bwd_ostride = stride_api_choice == dft::detail::stride_api::FB_STRIDES ? a_stride : b_stride; - if (a_min - stride_vecs.vec_a.begin() != rank) { + if (a_min_idx != rank) { // swap dimensions to have the last one have the smallest stride - std::swap(n_copy[a_min - stride_vecs.vec_a.begin() - 1], n_copy[rank - 1]); + std::swap(n_copy[a_min_idx - 1], n_copy[rank - 1]); } for (int i = 1; i < rank; i++) { if ((stride_vecs.vec_a[i] % a_stride != 0) || (stride_vecs.vec_b[i] % b_stride != 0)) {