Hi-PACE / hipace

Highly efficient Plasma Accelerator Emulation, quasistatic particle-in-cell code
https://hipace.readthedocs.io
Other
51 stars 14 forks source link

Refactor AnyFFT #1099

Closed AlexanderSinn closed 4 months ago

AlexanderSinn commented 4 months ago

In this PR, the functions that calculate a discrete sine transformation using an FFT are moved from AnyDST to the Poisson solvers.

This PR also fixes FFTDirichletExpanded for rocfft by setting the expanded position array (the FFT input) to zero after every FFT. Presumably, rocfft is using the array as work area. 

New AnyFFT interface:

struct VendorPlan;

enum struct FFTType {
    C2C_2D_fwd,
    C2C_2D_bkw,
    C2R_2D,
    R2C_2D,
    R2R_2D,
    C2R_1D_batched
};

struct AnyFFT {

    /** \brief Initialize an FFT plan for the requested transform type using a Vendor FFT library.
     * For 1D batched transforms, ny represents the number of batches to calculate at once.
     * This function returns the number of bytes of the work area needed for the FFT. The work area
     * has to be allocated by the function that uses the FFT and passed into SetBuffers.
     *
     * \param[in] type Type of FFT to perform
     * \param[in] nx Size of the contiguous dimension of the FFT
     * \param[in] ny Size of the second dimension of the FFT
     */
    std::size_t Initialize (FFTType type, int nx, int ny);

    /** \brief Set the pointers to the input, output and work area of the FFT.
     * This function has to be called after Initialize and before Execute.
     *
     * \param[in] in Pointer to the input of the FFT
     * \param[in] out Pointer to the output of the FFT
     * \param[in] work_area Pointer to the work area for the FFT
     */
    void SetBuffers (void* in, void* out, void* work_area);

    /** \brief Perform the initialized FFT */
    void Execute ();

    /** \brief Destructor to destroy the FFT plan */
    ~AnyFFT ();

    /** \brief Setup function that has to be called before any FFT plan is initialized. */
    static void setup ();

    /** \brief Cleanup function that has to be called at the end of the program. */
    static void cleanup ();

private:
    /** Vendor specific data for the FFT */
    VendorPlan* m_plan = nullptr;
};

Testing: