Implement progressive DFT execution

dancazarin · dancazarin · commit b6d4e797fe13 · 2025-06-15T07:06:10.000+02:00
diff --git a/include/kfr/dft/fft.hpp b/include/kfr/dft/fft.hpp
@@ -92,6 +92,13 @@ struct dft_stage
     {
         do_execute(cinvert_t(), out, in, temp);
     }
+    KFR_MEM_INTRINSIC void execute(bool inverse, complex<T>* out, const complex<T>* in, u8* temp)
+    {
+        if (inverse)
+            do_execute(cinvert_t(), out, in, temp);
+        else
+            do_execute(cdirect_t(), out, in, temp);
+    }
     virtual ~dft_stage() {}
 
 protected:
@@ -155,6 +162,12 @@ using fn_transpose = void (*)(complex<T>*, const complex<T>*, shape<2>);
 template <typename T>
 void dft_initialize_transpose(fn_transpose<T>& transpose);
 
+template <typename T>
+void dft_progressive_start(const dft_plan<T>& plan, typename dft_plan<T>::progressive& progressive,
+                           bool inverse, complex<T>* out, const complex<T>* in, u8* temp);
+template <typename T>
+void dft_progressive_step(const dft_plan<T>& plan, typename dft_plan<T>::progressive& progressive);
+
 } // namespace internal_generic
 
 /**
@@ -232,9 +245,10 @@ struct dft_plan
      *
      * @param size The size of the DFT.
      * @param order The order of the DFT samples. See `dft_order`.
+     * @param progressive_optimized If true, the plan will be optimized for progressive execution.
      */
-    explicit dft_plan(size_t size, dft_order order = dft_order::normal)
-        : size(size), temp_size(0), data_size(0), arblen(false)
+    explicit dft_plan(size_t size, dft_order order = dft_order::normal, bool progressive_optimized = false)
+        : size(size), temp_size(0), data_size(0), arblen(false), progressive_optimized(progressive_optimized)
     {
         internal_generic::dft_initialize(*this);
     }
@@ -369,6 +383,7 @@ struct dft_plan
     std::vector<dft_stage_ptr<T>> all_stages; /**< Internal data. */
     std::array<std::vector<dft_stage<T>*>, 2> stages; /**< Internal data. */
     bool arblen; /**< True if Bluestein's FFT algorithm is selected. */
+    bool progressive_optimized; /**< True if the plan is for progressive execution of the DFT. */
     using bitset = std::bitset<DFT_MAX_STAGES>; /**< Internal typedef. */
     std::array<bitset, 2> disposition_inplace; /**< Internal data. */
     std::array<bitset, 2> disposition_outofplace; /**< Internal data. */
@@ -380,12 +395,64 @@ struct dft_plan
     static bitset precompute_disposition(int num_stages, bitset can_inplace_per_stage,
                                          bool inplace_requested);
 
+    /** Internal data structure for progressive execution of the DFT.
+        Do not access the members directly as they may change in future versions.
+     */
+    struct progressive
+    {
+        bool inverse;
+        complex<T>* out;
+        const complex<T>* in;
+        u8* temp;
+        bitset disposition;
+        complex<T>* scratch;
+        size_t step = 0;
+    };
+
+    /// @brief Returns the number of steps for progressive execution of the DFT.
+    /// @return The number of steps for progressive execution.
+    size_t progressive_total_steps() const;
+
+    /**
+     * @brief Initiates the progressive execution of the DFT.
+     * @param inverse If true, applies the inverse DFT.
+     * @param out Pointer to the output data.
+     * @param in Pointer to the input data.
+     * @param temp Temporary (scratch) buffer. A scratch buffer of size
+     * `plan->temp_size` must be provided.
+     * @return A `progressive` structure that can be used with `progressive_step`.
+     * @note Ensure that the entire input data is available in the `in` buffer before calling this function.
+     * The `out` buffer will contain the result data after the final step of the progressive execution.
+     */
+    KFR_MEM_INTRINSIC progressive progressive_start(bool inverse, complex<T>* out, const complex<T>* in,
+                                                    u8* temp) const
+    {
+        KFR_LOGIC_CHECK(is_initialized(), "dft_plan is not initialized");
+        KFR_LOGIC_CHECK(temp_size == 0 || temp != nullptr,
+                        "Temporary buffer must be provided for progressive execution");
+        progressive result{};
+        internal_generic::dft_progressive_start(*this, result, inverse, out, in, temp);
+        return result;
+    }
+
+    /**
+     * @brief Steps the progressive execution of the DFT.
+     * @param progressive A `progressive` structure returned by `progressive_start`.
+     * @return `true` if there are more steps to execute, `false` if the DFT is complete.
+     */
+    KFR_MEM_INTRINSIC bool progressive_step(progressive& progressive) const
+    {
+        internal_generic::dft_progressive_step(*this, progressive);
+        return ++progressive.step < stages[progressive.inverse].size();
+    }
+
 protected:
     struct noinit
     {
     };
-    explicit dft_plan(noinit, size_t size, dft_order order = dft_order::normal)
-        : size(size), temp_size(0), data_size(0), arblen(false)
+    explicit dft_plan(noinit, size_t size, dft_order order = dft_order::normal,
+                      bool progressive_optimized = false)
+        : size(size), temp_size(0), data_size(0), arblen(false), progressive_optimized(progressive_optimized)
     {
     }
 
@@ -426,8 +493,10 @@ struct dft_plan_real : dft_plan<T>
         (void)cpu;
     }
 
-    explicit dft_plan_real(size_t size, dft_pack_format fmt = dft_pack_format::CCs)
-        : dft_plan<T>(typename dft_plan<T>::noinit{}, size / 2), size(size), fmt(fmt)
+    explicit dft_plan_real(size_t size, dft_pack_format fmt = dft_pack_format::CCs,
+                           bool progressive_optimized = false)
+        : dft_plan<T>(typename dft_plan<T>::noinit{}, size / 2, dft_order::normal, progressive_optimized),
+          size(size), fmt(fmt)
     {
         KFR_LOGIC_CHECK(is_even(size), "dft_plan_real requires size to be even");
         internal_generic::dft_real_initialize(*this);
@@ -488,6 +557,28 @@ struct dft_plan_real : dft_plan<T>
     {
         this->execute_dft(ctrue, ptr_cast<complex<T>>(out.data()), in.data(), temp);
     }
+
+    using progressive = typename dft_plan<T>::progressive;
+
+    KFR_MEM_INTRINSIC progressive progressive_start(T* out, const complex<T>* in, u8* temp) const
+    {
+        KFR_LOGIC_CHECK(is_initialized(), "dft_plan_real is not initialized");
+        KFR_LOGIC_CHECK(this->temp_size == 0 || temp != nullptr,
+                        "Temporary buffer must be provided for progressive execution");
+        progressive result{};
+        internal_generic::dft_progressive_start(*this, result, true, ptr_cast<complex<T>>(out), in, temp);
+        return result;
+    }
+    KFR_MEM_INTRINSIC progressive progressive_start(complex<T>* out, const T* in, u8* temp) const
+    {
+        KFR_LOGIC_CHECK(is_initialized(), "dft_plan_real is not initialized");
+        KFR_LOGIC_CHECK(this->temp_size == 0 || temp != nullptr,
+                        "Temporary buffer must be provided for progressive execution");
+        progressive result{};
+        internal_generic::dft_progressive_start(*this, result, false, out, ptr_cast<const complex<T>>(in),
+                                                temp);
+        return result;
+    }
 };
 
 /// @brief Multidimensional DFT
diff --git a/src/dft/dft.cpp b/src/dft/dft.cpp
@@ -39,6 +39,12 @@ void dft_plan<T>::dump() const
     }
 }
 
+template <typename T>
+size_t dft_plan<T>::progressive_total_steps() const
+{
+    return stages[0].size();
+}
+
 template <typename T>
 void dft_plan<T>::calc_disposition()
 {
@@ -131,6 +137,13 @@ CMT_MULTI_PROTO(namespace impl {
                      u8* temp);
     template <typename T>
     void dft_initialize_transpose(internal_generic::fn_transpose<T> & transpose);
+
+    template <typename T>
+    void dft_progressive_start(const dft_plan<T>& plan, typename dft_plan<T>::progressive& progressive,
+                               bool inverse, complex<T>* out, const complex<T>* in, u8* temp);
+
+    template <typename T>
+    void dft_progressive_step(const dft_plan<T>& plan, typename dft_plan<T>::progressive& progressive);
 })
 
 #ifdef CMT_MULTI_NEEDS_GATE
@@ -159,6 +172,19 @@ void dft_initialize_transpose(fn_transpose<T>& transpose)
     CMT_MULTI_GATE(ns::impl::dft_initialize_transpose(transpose));
 }
 
+template <typename T>
+void dft_progressive_start(const dft_plan<T>& plan, typename dft_plan<T>::progressive& progressive,
+                           bool inverse, complex<T>* out, const complex<T>* in, u8* temp)
+{
+    CMT_MULTI_GATE(ns::impl::dft_progressive_start(plan, progressive, inverse, out, in, temp));
+}
+
+template <typename T>
+void dft_progressive_step(const dft_plan<T>& plan, typename dft_plan<T>::progressive& progressive)
+{
+    CMT_MULTI_GATE(ns::impl::dft_progressive_step(plan, progressive));
+}
+
 template void dft_initialize<float>(dft_plan<float>&);
 template void dft_initialize<double>(dft_plan<double>&);
 template void dft_real_initialize<float>(dft_plan_real<float>&);
@@ -173,6 +199,16 @@ template void dft_execute<double>(const dft_plan<double>&, cbool_t<true>, comple
                                   const complex<double>*, u8*);
 template void dft_initialize_transpose<float>(fn_transpose<float>&);
 template void dft_initialize_transpose<double>(fn_transpose<double>&);
+template void dft_progressive_start(const dft_plan<float>& plan,
+                                    typename dft_plan<float>::progressive& progressive, bool inverse,
+                                    complex<float>* out, const complex<float>* in, u8* temp);
+template void dft_progressive_start(const dft_plan<double>& plan,
+                                    typename dft_plan<double>::progressive& progressive, bool inverse,
+                                    complex<double>* out, const complex<double>* in, u8* temp);
+template void dft_progressive_step(const dft_plan<float>& plan,
+                                   typename dft_plan<float>::progressive& progressive);
+template void dft_progressive_step(const dft_plan<double>& plan,
+                                   typename dft_plan<double>::progressive& progressive);
 
 } // namespace internal_generic
 
diff --git a/src/dft/fft-impl.hpp b/src/dft/fft-impl.hpp
@@ -862,7 +862,6 @@ KFR_INTRINSIC ctrue_t radix4_pass(csize_t<4>, size_t blocks, csize_t<width>, cfa
 template <typename T>
 struct fft_config
 {
-    constexpr static inline const bool recursion = true;
 #ifdef CMT_ARCH_NEON
     constexpr static inline const bool prefetch = false;
 #else
@@ -881,7 +880,7 @@ struct fft_stage_impl : dft_stage<T>
         this->radix      = 4;
         this->stage_size = stage_size;
         this->repeats    = 4;
-        this->recursion  = fft_config<T>::recursion;
+        this->recursion  = true;
         this->data_size =
             align_up(sizeof(complex<T>) * stage_size / 4 * 3, platform<>::native_cache_alignment);
     }
@@ -922,7 +921,7 @@ struct fft_final_stage_impl : dft_stage<T>
         this->stage_size = size;
         this->out_offset = size;
         this->repeats    = 4;
-        this->recursion  = fft_config<T>::recursion;
+        this->recursion  = true;
         this->data_size  = align_up(sizeof(complex<T>) * size * 3 / 2, platform<>::native_cache_alignment);
     }
 
@@ -1706,9 +1705,9 @@ void make_fft_stages(dft_plan<T>* self, cbool_t<autosort>, size_t stage_size, cb
 } // namespace intrinsics
 
 template <bool is_even, typename T>
-void make_fft(dft_plan<T>* self, size_t stage_size, cbool_t<is_even>)
+void make_fft(dft_plan<T>* self, size_t stage_size, cbool_t<is_even>, bool autosort)
 {
-    if (use_autosort<T>(ilog2(stage_size)))
+    if (autosort)
     {
         intrinsics::make_fft_stages(self, ctrue, stage_size, cbool<is_even>, ctrue);
     }
@@ -1776,7 +1775,8 @@ KFR_INTRINSIC void initialize_order(dft_plan<T>* self)
 template <typename T>
 KFR_INTRINSIC void init_fft(dft_plan<T>* self, size_t size, dft_order)
 {
-    const size_t log2n = ilog2(size);
+    const size_t log2n  = ilog2(size);
+    const bool autosort = use_autosort<T>(ilog2(size)) || self->progressive_optimized;
     cswitch(
         csizes_t<0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10
 #ifdef KFR_AUTOSORT_FOR_2048
@@ -1791,8 +1791,10 @@ KFR_INTRINSIC void init_fft(dft_plan<T>* self, size_t size, dft_order)
             constexpr size_t log2nv = val_of(decltype(log2n)());
             add_stage<intrinsics::fft_specialization<T, log2nv>>(self, size);
         },
-        [&]()
-        { cswitch(cfalse_true, is_even(log2n), [&](auto is_even) { make_fft(self, size, is_even); }); });
+        [&]() {
+            cswitch(cfalse_true, is_even(log2n),
+                    [&](auto is_even) { make_fft(self, size, is_even, autosort); });
+        });
 }
 
 template <typename T>
@@ -2027,14 +2029,14 @@ void dft_execute(const dft_plan<T>& plan, cbool_t<inverse>, complex<T>* out, con
         }
         else
         {
-            size_t offset = 0;
+            size_t offset            = 0;
+            complex<T>* cur_out      = select_out(plan, disposition, depth, stages.size(), out, scratch);
+            const complex<T>* cur_in = select_in(plan, disposition, depth, out, in, scratch);
+            dft_stage<T>* stage      = stages[depth];
             while (offset < plan.size)
             {
-                stages[depth]->execute(cbool<inverse>,
-                                       select_out(plan, disposition, depth, stages.size(), out, scratch) +
-                                           offset,
-                                       select_in(plan, disposition, depth, out, in, scratch) + offset, temp);
-                offset += stages[depth]->stage_size;
+                stage->execute(cbool<inverse>, cur_out + offset, cur_in + offset, temp);
+                offset += stage->stage_size;
             }
             depth++;
         }
@@ -2045,6 +2047,47 @@ void dft_initialize_transpose(internal_generic::fn_transpose<T>& transpose)
 {
     transpose = &kfr::CMT_ARCH_NAME::matrix_transpose;
 }
+
+template <typename T>
+void dft_progressive_start(const dft_plan<T>& plan, typename dft_plan<T>::progressive& prog, bool inverse,
+                           complex<T>* out, const complex<T>* in, u8* temp)
+{
+    prog.inverse = inverse;
+    prog.out     = out;
+    prog.in      = in;
+    prog.temp    = temp;
+    prog.scratch = ptr_cast<complex<T>>(
+        temp + plan.temp_size -
+        align_up(sizeof(complex<T>) * (plan.size + 1), platform<>::native_cache_alignment));
+
+    prog.disposition = in == out ? plan.disposition_inplace[inverse] : plan.disposition_outofplace[inverse];
+
+    bool in_scratch = prog.disposition.test(0);
+    if (in_scratch)
+    {
+        plan.stages[inverse][0]->copy_input(inverse, prog.scratch, in, plan.size);
+    }
+    prog.step = 0;
+}
+
+template <typename T>
+void dft_progressive_step(const dft_plan<T>& plan, typename dft_plan<T>::progressive& progressive)
+{
+    auto&& stages  = plan.stages[progressive.inverse];
+    uint32_t depth = progressive.step;
+    complex<T>* cur_out =
+        select_out(plan, progressive.disposition, depth, stages.size(), progressive.out, progressive.scratch);
+    const complex<T>* cur_in =
+        select_in(plan, progressive.disposition, depth, progressive.out, progressive.in, progressive.scratch);
+
+    size_t offset       = 0;
+    dft_stage<T>* stage = stages[depth];
+    while (offset < plan.size)
+    {
+        stage->execute(progressive.inverse, cur_out + offset, cur_in + offset, progressive.temp);
+        offset += stage->stage_size;
+    }
+}
 } // namespace impl
 
 namespace intrinsics
diff --git a/src/dft/fft-templates.hpp b/src/dft/fft-templates.hpp
@@ -33,13 +33,16 @@ inline namespace CMT_ARCH_NAME
 {
 namespace impl
 {
-template void dft_initialize<FLOAT>(dft_plan<FLOAT>& plan);
-template void dft_real_initialize<FLOAT>(dft_plan_real<FLOAT>& plan);
+template void dft_initialize<FLOAT>(dft_plan<FLOAT>&);
+template void dft_real_initialize<FLOAT>(dft_plan_real<FLOAT>&);
 template void dft_execute<FLOAT>(const dft_plan<FLOAT>&, cbool_t<false>, complex<FLOAT>*,
                                  const complex<FLOAT>*, u8*);
 template void dft_execute<FLOAT>(const dft_plan<FLOAT>&, cbool_t<true>, complex<FLOAT>*,
                                  const complex<FLOAT>*, u8*);
-template void dft_initialize_transpose<FLOAT>(internal_generic::fn_transpose<FLOAT>& transpose);
+template void dft_initialize_transpose<FLOAT>(internal_generic::fn_transpose<FLOAT>&);
+template void dft_progressive_start(const dft_plan<FLOAT>&, typename dft_plan<FLOAT>::progressive&, bool,
+                                    complex<FLOAT>*, const complex<FLOAT>*, u8*);
+template void dft_progressive_step(const dft_plan<FLOAT>&, typename dft_plan<FLOAT>::progressive&);
 } // namespace impl
 } // namespace CMT_ARCH_NAME
 } // namespace kfr
diff --git a/tests/dft_test.cpp b/tests/dft_test.cpp
