Add multidimensional DFT to C API

kfrlib · Nov 18, 2024 · 95b21ad · 95b21ad
1 parent 5191a48
commit 95b21ad
Show file tree

Hide file tree

Showing 4 changed files with 683 additions and 377 deletions.
diff --git a/include/kfr/capi.h b/include/kfr/capi.h
@@ -184,13 +184,63 @@ typedef enum KFR_DFT_PACK_FORMAT
  */
 KFR_API_SPEC KFR_DFT_PLAN_F32* kfr_dft_create_plan_f32(size_t size);
 
+/**
+ * @brief Create a 2D complex DFT plan (Single precision).
+ * @param size1 Size of the first dimension.
+ * @param size2 Size of the second dimension.
+ * @return Pointer to the created 2D DFT plan. Use `kfr_dft_delete_plan_f**` to free.
+ */
+KFR_API_SPEC KFR_DFT_PLAN_F32* kfr_dft_create_2d_plan_f32(size_t size1, size_t size2);
+
+/**
+ * @brief Create a 3D complex DFT plan (Single precision).
+ * @param size1 Size of the first dimension.
+ * @param size2 Size of the second dimension.
+ * @param size3 Size of the third dimension.
+ * @return Pointer to the created 3D DFT plan. Use `kfr_dft_delete_plan_f**` to free.
+ */
+KFR_API_SPEC KFR_DFT_PLAN_F32* kfr_dft_create_3d_plan_f32(size_t size1, size_t size2, size_t size3);
+
+/**
+ * @brief Create an N-dimensional complex DFT plan (Single precision).
+ * @param dims Number of dimensions.
+ * @param shape Array of sizes for each dimension.
+ * @return Pointer to the created N-dimensional DFT plan. Use `kfr_dft_delete_plan_f**` to free.
+ */
+KFR_API_SPEC KFR_DFT_PLAN_F32* kfr_dft_create_md_plan_f32(size_t dims, const unsigned* shape);
+
 /**
  * @brief Create a complex DFT plan (Double precision).
  * @param size Size of the DFT.
  * @return Pointer to the created DFT plan. Use `kfr_dft_delete_plan_f**` to free.
  */
 KFR_API_SPEC KFR_DFT_PLAN_F64* kfr_dft_create_plan_f64(size_t size);
 
+/**
+ * @brief Create a 2D complex DFT plan (Double precision).
+ * @param size1 Size of the first dimension.
+ * @param size2 Size of the second dimension.
+ * @return Pointer to the created 2D DFT plan. Use `kfr_dft_delete_plan_f**` to free.
+ */
+KFR_API_SPEC KFR_DFT_PLAN_F64* kfr_dft_create_2d_plan_f64(size_t size1, size_t size2);
+
+/**
+ * @brief Create a 3D complex DFT plan (Double precision).
+ * @param size1 Size of the first dimension.
+ * @param size2 Size of the second dimension.
+ * @param size3 Size of the third dimension.
+ * @return Pointer to the created 3D DFT plan. Use `kfr_dft_delete_plan_f**` to free.
+ */
+KFR_API_SPEC KFR_DFT_PLAN_F64* kfr_dft_create_3d_plan_f64(size_t size1, size_t size2, size_t size3);
+
+/**
+ * @brief Create an N-dimensional complex DFT plan (Double precision).
+ * @param dims Number of dimensions.
+ * @param shape Array of sizes for each dimension.
+ * @return Pointer to the created N-dimensional DFT plan. Use `kfr_dft_delete_plan_f**` to free.
+ */
+KFR_API_SPEC KFR_DFT_PLAN_F64* kfr_dft_create_md_plan_f64(size_t dims, const unsigned* shape);
+
 /**
  * @brief Dump details of the DFT plan to stdout for inspection.
  * @param plan Pointer to the DFT plan.
@@ -308,6 +358,13 @@ KFR_API_SPEC void kfr_dft_delete_plan_f64(KFR_DFT_PLAN_F64* plan);
 KFR_API_SPEC KFR_DFT_REAL_PLAN_F32* kfr_dft_real_create_plan_f32(size_t size,
                                                                  KFR_DFT_PACK_FORMAT pack_format);
 
+KFR_API_SPEC KFR_DFT_REAL_PLAN_F32* kfr_dft_real_create_2d_plan_f32(size_t size1, size_t size2,
+                                                                    bool real_out_is_enough);
+KFR_API_SPEC KFR_DFT_REAL_PLAN_F32* kfr_dft_real_create_3d_plan_f32(size_t size1, size_t size2, size_t size3,
+                                                                    bool real_out_is_enough);
+KFR_API_SPEC KFR_DFT_REAL_PLAN_F32* kfr_dft_real_create_md_plan_f32(size_t dims, const unsigned* shape,
+                                                                    bool real_out_is_enough);
+
 /**
  * @brief Create a real DFT plan (Double precision).
  * @param size Size of the real DFT. Must be even.
@@ -317,6 +374,13 @@ KFR_API_SPEC KFR_DFT_REAL_PLAN_F32* kfr_dft_real_create_plan_f32(size_t size,
 KFR_API_SPEC KFR_DFT_REAL_PLAN_F64* kfr_dft_real_create_plan_f64(size_t size,
                                                                  KFR_DFT_PACK_FORMAT pack_format);
 
+KFR_API_SPEC KFR_DFT_REAL_PLAN_F64* kfr_dft_real_create_2d_plan_f64(size_t size1, size_t size2,
+                                                                    bool real_out_is_enough);
+KFR_API_SPEC KFR_DFT_REAL_PLAN_F64* kfr_dft_real_create_3d_plan_f64(size_t size1, size_t size2, size_t size3,
+                                                                    bool real_out_is_enough);
+KFR_API_SPEC KFR_DFT_REAL_PLAN_F64* kfr_dft_real_create_md_plan_f64(size_t dims, const unsigned* shape,
+                                                                    bool real_out_is_enough);
+
 /**
  * @brief Dump details of the real DFT plan to stdout for inspection.
  * @param plan Pointer to the DFT plan.

diff --git a/include/kfr/dft/fft.hpp b/include/kfr/dft/fft.hpp
@@ -354,7 +354,7 @@ struct dft_plan_real : dft_plan<T>
 };
 
 /// @brief Multidimensional DFT
-template <typename T, index_t Dims>
+template <typename T, index_t Dims = dynamic_shape>
 struct dft_plan_md
 {
     shape<Dims> size;
@@ -490,7 +490,7 @@ struct dft_plan_md
 };
 
 /// @brief Multidimensional DFT
-template <typename T, index_t Dims>
+template <typename T, index_t Dims = dynamic_shape>
 struct dft_plan_md_real
 {
     shape<Dims> size;