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)) {