luat_lib_fft.c 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886
  1. /*
  2. @module fft
  3. @summary 快速傅里叶变换(FFT/IFFT),支持 float32 与 q15 定点内核
  4. @version 0.2
  5. @date 2025.08
  6. @demo fft
  7. @tag LUAT_USE_FFT
  8. @usage
  9. -- 模块通常作为内置库,无需 require,直接调用
  10. -- 示例见各 API 注释以及仓库 demo/test 脚本
  11. -- 支持的点数范围:
  12. -- 输入点数N,要求为2的幂次方,且N < 65536
  13. -- 格式支持:Lua数组 或 zbuff输入
  14. -- 推荐范围:N ≤ 16384(已验证稳定运行)
  15. -- 性能测试结果(实测数据):
  16. -- Air780EHV平台:
  17. -- 2048点:F32=23ms, Q15=10ms(Q15快2.3倍)
  18. -- 16384点:F32=240ms, Q15=110ms(Q15快2.2倍)
  19. -- Air8101平台(有硬件浮点加速):
  20. -- 2048点:F32=6ms, Q15=5ms(性能相近)
  21. -- 16384点:F32=78ms, Q15=86ms(F32略快)
  22. -- 平台建议:
  23. -- Air780EHV/Air780EPM:推荐使用Q15路径,性能提升显著
  24. -- Air8101:小规模FFT两种都可,大规模FFT建议F32路径
  25. -- 内存受限场景:优先选择Q15路径(内存需求减半)
  26. -- 高精度场景:优先选择F32路径
  27. */
  28. #include "luat_base.h"
  29. #include "luat_mem.h"
  30. #include <math.h>
  31. #include <string.h>
  32. #include "fft_core.h"
  33. #define LUAT_LOG_TAG "fft"
  34. #include "luat_log.h"
  35. #include "luat_conf_bsp.h"
  36. #include "luat_zbuff.h"
  37. #include "rotable2.h"
  38. // Q15 内核头文件(当前版本已实现)
  39. #include "fft_core_q15.h"
  40. // helper: read float array from lua table (1-based)
  41. static int read_lua_array_float(lua_State* L, int idx, float* out, int n)
  42. {
  43. for (int i = 0; i < n; i++) {
  44. lua_rawgeti(L, idx, i + 1);
  45. if (!lua_isnumber(L, -1)) {
  46. lua_pop(L, 1);
  47. return -1;
  48. }
  49. out[i] = (float)lua_tonumber(L, -1);
  50. lua_pop(L, 1);
  51. }
  52. return 0;
  53. }
  54. // helper: write float array into lua table (1-based)
  55. static void write_lua_array_float(lua_State* L, float* a, int n)
  56. {
  57. lua_createtable(L, n, 0);
  58. for (int i = 0; i < n; i++) {
  59. lua_pushnumber(L, a[i]);
  60. lua_rawseti(L, -2, i + 1);
  61. }
  62. }
  63. /*
  64. 生成 float32 旋转因子
  65. @api fft.generate_twiddles(N)
  66. @int N 点数,必须为 2 的幂
  67. @return table Wc, table Ws 两个 Lua 数组(长度 N/2),分别为 cos 与 -sin
  68. @usage
  69. local N = 2048
  70. local Wc, Ws = fft.generate_twiddles(N)
  71. */
  72. static int l_fft_generate_twiddles(lua_State* L)
  73. {
  74. int N = luaL_checkinteger(L, 1);
  75. if (N <= 1 || (N & (N - 1)))
  76. return luaL_error(L, "N must be power of 2");
  77. int half = N / 2;
  78. float* Wc = luat_heap_malloc(sizeof(float) * half);
  79. float* Ws = luat_heap_malloc(sizeof(float) * half);
  80. if (!Wc || !Ws) {
  81. if (Wc)
  82. luat_heap_free(Wc);
  83. if (Ws)
  84. luat_heap_free(Ws);
  85. return luaL_error(L, "no memory");
  86. }
  87. luat_fft_generate_twiddles(N, Wc, Ws);
  88. write_lua_array_float(L, Wc, half);
  89. write_lua_array_float(L, Ws, half);
  90. luat_heap_free(Wc);
  91. luat_heap_free(Ws);
  92. return 2;
  93. }
  94. // 生成 Q15 旋转因子到 zbuff(长度要求:N/2 * 2 字节)
  95. // 调用方式:fft.generate_twiddles_q15_to_zbuff(N, Wc_zbuff, Ws_zbuff)
  96. /*
  97. 生成 q15 旋转因子到 zbuff(零浮点)
  98. @api fft.generate_twiddles_q15_to_zbuff(N, Wc_zb, Ws_zb)
  99. @int N 点数,必须为 2 的幂
  100. @zbuff Wc_zb 输出缓冲,长度至少为 (N/2)*2 字节,存放 int16 Q15 的 cos
  101. @zbuff Ws_zb 输出缓冲,长度至少为 (N/2)*2 字节,存放 int16 Q15 的 -sin(前向)
  102. @return nil 无返回值,结果写入传入的 zbuff
  103. @usage
  104. local N = 2048
  105. local Wc_q15 = zbuff.create((N//2)*2)
  106. local Ws_q15 = zbuff.create((N//2)*2)
  107. fft.generate_twiddles_q15_to_zbuff(N, Wc_q15, Ws_q15)
  108. */
  109. static int l_fft_generate_twiddles_q15_to_zbuff(lua_State* L)
  110. {
  111. // 使用整型查表(度数,1度分辨率,缩放256)生成近似Q15旋转因子
  112. static const int16_t SIN_TABLE256[91] = {
  113. 0, 4, 9, 13, 18, 22, 27, 31, 36, 40, 44, 49, 53, 58, 62, 66, 71, 75, 79, 83,
  114. 88, 92, 96, 100, 104, 108, 112, 116, 120, 124, 128, 132, 136, 139, 143, 147, 150, 154, 158, 161,
  115. 165, 168, 171, 175, 178, 181, 184, 187, 190, 193, 196, 199, 202, 204, 207, 210, 212, 215, 217, 219,
  116. 222, 224, 226, 228, 230, 232, 234, 236, 237, 239, 241, 242, 243, 245, 246, 247, 248, 249, 250, 251,
  117. 252, 253, 254, 254, 255, 255, 255, 256, 256, 256, 256
  118. };
  119. int N = luaL_checkinteger(L, 1);
  120. if (N <= 1 || (N & (N - 1)))
  121. return luaL_error(L, "N 必须为 2 的幂");
  122. luat_zbuff_t* zbWc = (luat_zbuff_t*)luaL_testudata(L, 2, LUAT_ZBUFF_TYPE);
  123. luat_zbuff_t* zbWs = (luat_zbuff_t*)luaL_testudata(L, 3, LUAT_ZBUFF_TYPE);
  124. if (!zbWc || !zbWs)
  125. return luaL_error(L, "Wc/Ws 需为 zbuff");
  126. int need = (N / 2) * 2;
  127. if ((int)zbWc->len < need || (int)zbWs->len < need)
  128. return luaL_error(L, "zbuff 太小");
  129. int16_t* Wc = (int16_t*)zbWc->addr;
  130. int16_t* Ws = (int16_t*)zbWs->addr;
  131. for (int k = 0; k < N / 2; k++) {
  132. // angle = 360 * k / N (度),四舍五入
  133. int deg = (int)((int64_t)k * 360 + (N / 2)) / N;
  134. // cos 与 -sin,放大到 Q15(256<<7=32768,钳到 32767)
  135. int s256;
  136. int d = deg;
  137. int neg = 0;
  138. if (d < 0) {
  139. d = -d + 180;
  140. }
  141. d %= 360;
  142. if (d >= 180) {
  143. d -= 180;
  144. neg = 1;
  145. }
  146. if (d <= 90)
  147. s256 = SIN_TABLE256[d];
  148. else
  149. s256 = SIN_TABLE256[180 - d];
  150. int c256 = 0; // cos = sin(deg+90)
  151. int d2 = deg + 90;
  152. neg = 0;
  153. d = d2;
  154. if (d < 0) {
  155. d = -d + 180;
  156. }
  157. d %= 360;
  158. if (d >= 180) {
  159. d -= 180;
  160. neg = 1;
  161. }
  162. if (d <= 90)
  163. c256 = SIN_TABLE256[d];
  164. else
  165. c256 = SIN_TABLE256[180 - d];
  166. if (neg)
  167. c256 = -c256;
  168. int32_t wc = (int32_t)c256 << 7;
  169. if (wc > 32767)
  170. wc = 32767;
  171. if (wc < -32768)
  172. wc = -32768;
  173. // -sin:复用 s256 并加符号
  174. int sgn = 0;
  175. d = deg;
  176. if (d < 0) {
  177. d = -d + 180;
  178. }
  179. d %= 360;
  180. if (d >= 180) {
  181. d -= 180;
  182. sgn = 1;
  183. }
  184. int32_t ws = (int32_t)(sgn ? -s256 : s256) << 7;
  185. if (ws > 32767)
  186. ws = 32767;
  187. if (ws < -32768)
  188. ws = -32768;
  189. Wc[k] = (int16_t)wc;
  190. Ws[k] = (int16_t)ws;
  191. }
  192. return 0;
  193. }
  194. /*
  195. 原地 FFT 计算
  196. @api fft.run(real, imag, N, Wc, Ws[, opts])
  197. @param real 实部容器:
  198. - float32 路径:Lua 数组或 zbuff(float32)
  199. - q15 路径:zbuff(int16)。当 opts.core="q15" 且 opts.input_format 为 "u12"/"u16"/"s16" 时生效
  200. @param imag 虚部容器:同 real。可为 nil(视为全 0)
  201. @int N 点数,2 的幂
  202. @param Wc 旋转因子 cos:
  203. - float32 路径:Lua 数组或 zbuff(float32)
  204. - q15 路径:zbuff(int16),推荐配合 fft.generate_twiddles_q15_to_zbuff 生成
  205. @param Ws 旋转因子 -sin:同 Wc;IFFT 建议传入共轭版本以避免内层符号乘
  206. @table [opts]
  207. - core: "f32" | "q15"(默认 "f32")
  208. * "f32": 浮点内核,精度高(32位),计算稳定,适合精密分析
  209. * "q15": 定点内核,速度快(16位整数),内存省,适合实时处理但精度略低
  210. - input_format: "f32" | "u12" | "u16" | "s16"(q15 时必填其一)
  211. * "f32": 标准浮点输入,适用于已处理的信号数据
  212. * "u12": 12位无符号整数(0~4095),常见于ADC采样,自动去直流分量
  213. * "u16": 16位无符号整数(0~65535),适用于高精度ADC或预处理数据
  214. * "s16": 16位有符号整数(-32768~32767),适用于已去直流的差分信号
  215. @return nil 就地修改 real/imag
  216. @usage
  217. -- f32 路径示例(zbuff float32)
  218. local N=2048
  219. local real=zbuff.create(N*4); local imag=zbuff.create(N*4)
  220. local Wc,Ws=fft.generate_twiddles(N)
  221. fft.run(real, imag, N, Wc, Ws)
  222. -- q15 路径示例(U12 整数输入)
  223. local N=2048
  224. local real_i16=zbuff.create(N*2); local imag_i16=zbuff.create(N*2)
  225. local Wc_q15=zbuff.create((N//2)*2); local Ws_q15=zbuff.create((N//2)*2)
  226. fft.generate_twiddles_q15_to_zbuff(N, Wc_q15, Ws_q15)
  227. -- 写入 U12 数据到 real_i16 后:
  228. fft.run(real_i16, imag_i16, N, Wc_q15, Ws_q15, {core="q15", input_format="u12"})
  229. */
  230. static int l_fft_run(lua_State* L)
  231. {
  232. int N = luaL_checkinteger(L, 3);
  233. if (N <= 1 || (N & (N - 1)))
  234. return luaL_error(L, "N 必须为 2 的幂");
  235. float *r = NULL, *im = NULL, *Wc = NULL, *Ws = NULL;
  236. int r_free = 0, im_free = 0, wc_free = 0, ws_free = 0;
  237. // 可选参数解析(opts):当前版本仅支持 core/input_format(其余项作为未来优化)
  238. const char* core = "f32"; // "f32" | "q15"
  239. const char* input_format = "f32"; // "f32"|"u12"|"u16"|"s16"
  240. // real
  241. luat_zbuff_t* zb = (luat_zbuff_t*)luaL_testudata(L, 1, LUAT_ZBUFF_TYPE);
  242. if (zb) {
  243. r = (float*)zb->addr;
  244. } else {
  245. r = luat_heap_malloc(sizeof(float) * N);
  246. r_free = 1;
  247. if (!r)
  248. return luaL_error(L, "no memory");
  249. if (read_lua_array_float(L, 1, r, N)) {
  250. if (r_free)
  251. luat_heap_free(r);
  252. return luaL_error(L, "real must be number array or zbuff");
  253. }
  254. }
  255. // imag
  256. zb = (luat_zbuff_t*)luaL_testudata(L, 2, LUAT_ZBUFF_TYPE);
  257. if (zb) {
  258. im = (float*)zb->addr;
  259. } else {
  260. im = luat_heap_malloc(sizeof(float) * N);
  261. im_free = 1;
  262. if (!im) {
  263. if (r_free)
  264. luat_heap_free(r);
  265. return luaL_error(L, "no memory");
  266. }
  267. if (read_lua_array_float(L, 2, im, N)) {
  268. if (r_free)
  269. luat_heap_free(r);
  270. if (im_free)
  271. luat_heap_free(im);
  272. return luaL_error(L, "imag must be number array or zbuff");
  273. }
  274. }
  275. // W_real
  276. zb = (luat_zbuff_t*)luaL_testudata(L, 4, LUAT_ZBUFF_TYPE);
  277. if (zb) {
  278. Wc = (float*)zb->addr;
  279. } else {
  280. Wc = luat_heap_malloc(sizeof(float) * (N / 2));
  281. wc_free = 1;
  282. if (!Wc) {
  283. if (r_free)
  284. luat_heap_free(r);
  285. if (im_free)
  286. luat_heap_free(im);
  287. return luaL_error(L, "no memory");
  288. }
  289. if (read_lua_array_float(L, 4, Wc, N / 2)) {
  290. if (r_free)
  291. luat_heap_free(r);
  292. if (im_free)
  293. luat_heap_free(im);
  294. if (wc_free)
  295. luat_heap_free(Wc);
  296. return luaL_error(L, "W_real must be number array or zbuff");
  297. }
  298. }
  299. // W_imag
  300. zb = (luat_zbuff_t*)luaL_testudata(L, 5, LUAT_ZBUFF_TYPE);
  301. if (zb) {
  302. Ws = (float*)zb->addr;
  303. } else {
  304. Ws = luat_heap_malloc(sizeof(float) * (N / 2));
  305. ws_free = 1;
  306. if (!Ws) {
  307. if (r_free)
  308. luat_heap_free(r);
  309. if (im_free)
  310. luat_heap_free(im);
  311. if (wc_free)
  312. luat_heap_free(Wc);
  313. return luaL_error(L, "内存不足");
  314. }
  315. if (read_lua_array_float(L, 5, Ws, N / 2)) {
  316. if (r_free)
  317. luat_heap_free(r);
  318. if (im_free)
  319. luat_heap_free(im);
  320. if (wc_free)
  321. luat_heap_free(Wc);
  322. if (ws_free)
  323. luat_heap_free(Ws);
  324. return luaL_error(L, "W_imag 需为数字数组或 zbuff");
  325. }
  326. }
  327. // 读取第6个参数的 opts(若有)
  328. if (lua_gettop(L) >= 6 && lua_istable(L, 6)) {
  329. lua_getfield(L, 6, "core");
  330. if (!lua_isnil(L, -1))
  331. core = luaL_checkstring(L, -1);
  332. lua_pop(L, 1);
  333. lua_getfield(L, 6, "input_format");
  334. if (!lua_isnil(L, -1))
  335. input_format = luaL_checkstring(L, -1);
  336. lua_pop(L, 1);
  337. }
  338. // 如果选择 q15 内核,且输入为整数 zbuff,则走 q15 路径
  339. int use_q15 = (core && strcmp(core, "q15") == 0);
  340. int integer_input = (strcmp(input_format, "u12") == 0 || strcmp(input_format, "u16") == 0 || strcmp(input_format, "s16") == 0);
  341. if (use_q15 && integer_input) {
  342. // 校验 real/imag 是否为 zbuff(当前整型快速路径仅支持 zbuff)
  343. luat_zbuff_t* zb_real = (luat_zbuff_t*)luaL_testudata(L, 1, LUAT_ZBUFF_TYPE);
  344. luat_zbuff_t* zb_imag = (luat_zbuff_t*)luaL_testudata(L, 2, LUAT_ZBUFF_TYPE);
  345. if (!zb_real) {
  346. if (r_free)
  347. luat_heap_free(r);
  348. if (im_free)
  349. luat_heap_free(im);
  350. if (wc_free)
  351. luat_heap_free(Wc);
  352. if (ws_free)
  353. luat_heap_free(Ws);
  354. return luaL_error(L, "q15 模式要求 real 为整数 zbuff");
  355. }
  356. if ((int)zb_real->len < N * 2) {
  357. if (r_free)
  358. luat_heap_free(r);
  359. if (im_free)
  360. luat_heap_free(im);
  361. if (wc_free)
  362. luat_heap_free(Wc);
  363. if (ws_free)
  364. luat_heap_free(Ws);
  365. return luaL_error(L, "real zbuff 太小");
  366. }
  367. // 原地:将整数输入就地转换为带符号 Q15(覆盖 zbuff)
  368. uint16_t* r16 = (uint16_t*)zb_real->addr;
  369. for (int i = 0; i < N; i++) {
  370. int32_t v;
  371. uint16_t u = r16[i];
  372. if (strcmp(input_format, "u12") == 0) {
  373. v = (int32_t)(u & 0x0FFF) - 2048;
  374. v <<= 4;
  375. } else if (strcmp(input_format, "u16") == 0) {
  376. v = (int32_t)u - 32768;
  377. } else {
  378. v = (int16_t)u;
  379. }
  380. if (v > 32767)
  381. v = 32767;
  382. if (v < -32768)
  383. v = -32768;
  384. ((int16_t*)zb_real->addr)[i] = (int16_t)v;
  385. }
  386. if (zb_imag && (int)zb_imag->len >= N * 2) {
  387. uint16_t* i16 = (uint16_t*)zb_imag->addr;
  388. for (int i = 0; i < N; i++) {
  389. int32_t v;
  390. uint16_t u = i16[i];
  391. if (strcmp(input_format, "u12") == 0) {
  392. v = (int32_t)(u & 0x0FFF) - 2048;
  393. v <<= 4;
  394. } else if (strcmp(input_format, "u16") == 0) {
  395. v = (int32_t)u - 32768;
  396. } else {
  397. v = (int16_t)u;
  398. }
  399. if (v > 32767)
  400. v = 32767;
  401. if (v < -32768)
  402. v = -32768;
  403. ((int16_t*)zb_imag->addr)[i] = (int16_t)v;
  404. }
  405. } else if (zb_imag) {
  406. // 长度不足则清零
  407. memset(zb_imag->addr, 0, zb_imag->len);
  408. }
  409. // 强制要求外部传入 Q15 twiddle
  410. luat_zbuff_t* zbWc = (luat_zbuff_t*)luaL_testudata(L, 4, LUAT_ZBUFF_TYPE);
  411. luat_zbuff_t* zbWs = (luat_zbuff_t*)luaL_testudata(L, 5, LUAT_ZBUFF_TYPE);
  412. const int need_tw = (N / 2) * 2;
  413. if (!(zbWc && zbWs) || (int)zbWc->len < need_tw || (int)zbWs->len < need_tw) {
  414. if (r_free)
  415. luat_heap_free(r);
  416. if (im_free)
  417. luat_heap_free(im);
  418. if (wc_free)
  419. luat_heap_free(Wc);
  420. if (ws_free)
  421. luat_heap_free(Ws);
  422. return luaL_error(L, "q15 需传入 Wc/Ws zbuff,长度 N/2*2 字节");
  423. }
  424. int scale_exp = 0;
  425. int rc = luat_fft_inplace_q15((int16_t*)zb_real->addr, zb_imag ? (int16_t*)zb_imag->addr : NULL,
  426. N, 0, (const int16_t*)zbWc->addr, (const int16_t*)zbWs->addr,
  427. 0, &scale_exp);
  428. if (r_free)
  429. luat_heap_free(r);
  430. if (im_free)
  431. luat_heap_free(im);
  432. if (wc_free)
  433. luat_heap_free(Wc);
  434. if (ws_free)
  435. luat_heap_free(Ws);
  436. if (rc != 0)
  437. return luaL_error(L, "q15 内核执行失败");
  438. return 0;
  439. }
  440. // 默认:沿用 float32 路径
  441. luat_fft_run_inplace(r, im, N, Wc, Ws);
  442. // if input was table, write back
  443. if (!luaL_testudata(L, 1, LUAT_ZBUFF_TYPE)) {
  444. lua_settop(L, 2);
  445. for (int i = 0; i < N; i++) {
  446. lua_pushnumber(L, r[i]);
  447. lua_rawseti(L, 1, i + 1);
  448. }
  449. }
  450. if (!luaL_testudata(L, 2, LUAT_ZBUFF_TYPE)) {
  451. for (int i = 0; i < N; i++) {
  452. lua_pushnumber(L, im[i]);
  453. lua_rawseti(L, 2, i + 1);
  454. }
  455. }
  456. if (r_free)
  457. luat_heap_free(r);
  458. if (im_free)
  459. luat_heap_free(im);
  460. if (wc_free)
  461. luat_heap_free(Wc);
  462. if (ws_free)
  463. luat_heap_free(Ws);
  464. return 0;
  465. }
  466. /*
  467. 原地 IFFT 计算
  468. @api fft.ifft(real, imag, N, Wc, Ws[, opts])
  469. @param real 实部容器,类型与 fft.run 一致
  470. @param imag 虚部容器,类型与 fft.run 一致
  471. @int N 点数,2 的幂
  472. @param Wc 旋转因子 cos:类型同 fft.run
  473. @param Ws 旋转因子 -sin:类型同 fft.run。建议为 IFFT 预共轭传入 +sin 表
  474. @table [opts]
  475. - core: "f32" | "q15"(默认 "f32")
  476. - input_format: "f32" | "u12" | "u16" | "s16"(q15 时必填其一)
  477. @return nil 就地修改 real/imag,并在 f32 路径下包含 1/N 归一化
  478. */
  479. static int l_fft_ifft(lua_State* L)
  480. {
  481. int N = luaL_checkinteger(L, 3);
  482. if (N <= 1 || (N & (N - 1)))
  483. return luaL_error(L, "N 必须为 2 的幂");
  484. float *r = NULL, *im = NULL, *Wc = NULL, *Ws = NULL;
  485. int r_free = 0, im_free = 0, wc_free = 0, ws_free = 0;
  486. // 可选 opts(同 run):当前仅 core/input_format
  487. const char* core = "f32";
  488. const char* input_format = "f32";
  489. luat_zbuff_t* zb = NULL;
  490. zb = (luat_zbuff_t*)luaL_testudata(L, 1, LUAT_ZBUFF_TYPE);
  491. if (zb) {
  492. r = (float*)zb->addr;
  493. } else {
  494. r = luat_heap_malloc(sizeof(float) * N);
  495. r_free = 1;
  496. if (!r)
  497. return luaL_error(L, "no memory");
  498. if (read_lua_array_float(L, 1, r, N)) {
  499. if (r_free)
  500. luat_heap_free(r);
  501. return luaL_error(L, "real must be number array or zbuff");
  502. }
  503. }
  504. zb = (luat_zbuff_t*)luaL_testudata(L, 2, LUAT_ZBUFF_TYPE);
  505. if (zb) {
  506. im = (float*)zb->addr;
  507. } else {
  508. im = luat_heap_malloc(sizeof(float) * N);
  509. im_free = 1;
  510. if (!im) {
  511. if (r_free)
  512. luat_heap_free(r);
  513. return luaL_error(L, "no memory");
  514. }
  515. if (read_lua_array_float(L, 2, im, N)) {
  516. if (r_free)
  517. luat_heap_free(r);
  518. if (im_free)
  519. luat_heap_free(im);
  520. return luaL_error(L, "imag must be number array or zbuff");
  521. }
  522. }
  523. zb = (luat_zbuff_t*)luaL_testudata(L, 4, LUAT_ZBUFF_TYPE);
  524. if (zb) {
  525. Wc = (float*)zb->addr;
  526. } else {
  527. Wc = luat_heap_malloc(sizeof(float) * (N / 2));
  528. wc_free = 1;
  529. if (!Wc) {
  530. if (r_free)
  531. luat_heap_free(r);
  532. if (im_free)
  533. luat_heap_free(im);
  534. return luaL_error(L, "no memory");
  535. }
  536. if (read_lua_array_float(L, 4, Wc, N / 2)) {
  537. if (r_free)
  538. luat_heap_free(r);
  539. if (im_free)
  540. luat_heap_free(im);
  541. if (wc_free)
  542. luat_heap_free(Wc);
  543. return luaL_error(L, "W_real must be number array or zbuff");
  544. }
  545. }
  546. zb = (luat_zbuff_t*)luaL_testudata(L, 5, LUAT_ZBUFF_TYPE);
  547. if (zb) {
  548. Ws = (float*)zb->addr;
  549. } else {
  550. Ws = luat_heap_malloc(sizeof(float) * (N / 2));
  551. ws_free = 1;
  552. if (!Ws) {
  553. if (r_free)
  554. luat_heap_free(r);
  555. if (im_free)
  556. luat_heap_free(im);
  557. if (wc_free)
  558. luat_heap_free(Wc);
  559. return luaL_error(L, "内存不足");
  560. }
  561. if (read_lua_array_float(L, 5, Ws, N / 2)) {
  562. if (r_free)
  563. luat_heap_free(r);
  564. if (im_free)
  565. luat_heap_free(im);
  566. if (wc_free)
  567. luat_heap_free(Wc);
  568. if (ws_free)
  569. luat_heap_free(Ws);
  570. return luaL_error(L, "W_imag 需为数字数组或 zbuff");
  571. }
  572. }
  573. if (lua_gettop(L) >= 6 && lua_istable(L, 6)) {
  574. lua_getfield(L, 6, "core");
  575. if (!lua_isnil(L, -1))
  576. core = luaL_checkstring(L, -1);
  577. lua_pop(L, 1);
  578. lua_getfield(L, 6, "input_format");
  579. if (!lua_isnil(L, -1))
  580. input_format = luaL_checkstring(L, -1);
  581. lua_pop(L, 1);
  582. }
  583. int use_q15 = (core && strcmp(core, "q15") == 0);
  584. int integer_input = (strcmp(input_format, "u12") == 0 || strcmp(input_format, "u16") == 0 || strcmp(input_format, "s16") == 0);
  585. if (use_q15 && integer_input) {
  586. luat_zbuff_t* zb_real = (luat_zbuff_t*)luaL_testudata(L, 1, LUAT_ZBUFF_TYPE);
  587. luat_zbuff_t* zb_imag = (luat_zbuff_t*)luaL_testudata(L, 2, LUAT_ZBUFF_TYPE);
  588. if (!zb_real) {
  589. if (r_free)
  590. luat_heap_free(r);
  591. if (im_free)
  592. luat_heap_free(im);
  593. if (wc_free)
  594. luat_heap_free(Wc);
  595. if (ws_free)
  596. luat_heap_free(Ws);
  597. return luaL_error(L, "q15 模式要求 real 为整数 zbuff");
  598. }
  599. if ((int)zb_real->len < N * 2) {
  600. if (r_free)
  601. luat_heap_free(r);
  602. if (im_free)
  603. luat_heap_free(im);
  604. if (wc_free)
  605. luat_heap_free(Wc);
  606. if (ws_free)
  607. luat_heap_free(Ws);
  608. return luaL_error(L, "real zbuff 太小");
  609. }
  610. int16_t* rq = (int16_t*)luat_heap_malloc(sizeof(int16_t) * N);
  611. int16_t* iq = (int16_t*)luat_heap_malloc(sizeof(int16_t) * N);
  612. if (!rq || !iq) {
  613. if (rq)
  614. luat_heap_free(rq);
  615. if (iq)
  616. luat_heap_free(iq);
  617. if (r_free)
  618. luat_heap_free(r);
  619. if (im_free)
  620. luat_heap_free(im);
  621. if (wc_free)
  622. luat_heap_free(Wc);
  623. if (ws_free)
  624. luat_heap_free(Ws);
  625. return luaL_error(L, "内存不足");
  626. }
  627. const uint16_t* r16 = (const uint16_t*)zb_real->addr;
  628. for (int i = 0; i < N; i++) {
  629. int32_t v;
  630. uint16_t u = r16[i];
  631. if (strcmp(input_format, "u12") == 0) {
  632. v = ((int32_t)(u & 0x0FFF)) - 2048;
  633. v <<= 4;
  634. } else if (strcmp(input_format, "u16") == 0) {
  635. v = (int32_t)u - 32768;
  636. } else {
  637. v = (int16_t)u;
  638. }
  639. if (v > 32767)
  640. v = 32767;
  641. if (v < -32768)
  642. v = -32768;
  643. rq[i] = (int16_t)v;
  644. }
  645. if (zb_imag && (int)zb_imag->len >= N * 2) {
  646. const uint16_t* i16 = (const uint16_t*)zb_imag->addr;
  647. for (int i = 0; i < N; i++) {
  648. int32_t v;
  649. uint16_t u = i16[i];
  650. if (strcmp(input_format, "u12") == 0) {
  651. v = ((int32_t)(u & 0x0FFF)) - 2048;
  652. v <<= 4;
  653. } else if (strcmp(input_format, "u16") == 0) {
  654. v = (int32_t)u - 32768;
  655. } else {
  656. v = (int16_t)u;
  657. }
  658. if (v > 32767)
  659. v = 32767;
  660. if (v < -32768)
  661. v = -32768;
  662. iq[i] = (int16_t)v;
  663. }
  664. } else {
  665. memset(iq, 0, sizeof(int16_t) * N);
  666. }
  667. // 使用传入的 Q15 旋转因子(zbuff)或按需生成
  668. luat_zbuff_t* zbWc = (luat_zbuff_t*)luaL_testudata(L, 4, LUAT_ZBUFF_TYPE);
  669. luat_zbuff_t* zbWs = (luat_zbuff_t*)luaL_testudata(L, 5, LUAT_ZBUFF_TYPE);
  670. const int need_tw = (N / 2) * 2;
  671. const int16_t* Wcq = NULL;
  672. const int16_t* Wsq = NULL;
  673. int16_t* Wcq_alloc = NULL;
  674. int16_t* Wsq_alloc = NULL;
  675. if (zbWc && zbWs && (int)zbWc->len >= need_tw && (int)zbWs->len >= need_tw) {
  676. Wcq = (const int16_t*)zbWc->addr;
  677. Wsq = (const int16_t*)zbWs->addr;
  678. } else {
  679. Wcq_alloc = (int16_t*)luat_heap_malloc(sizeof(int16_t) * (N / 2));
  680. Wsq_alloc = (int16_t*)luat_heap_malloc(sizeof(int16_t) * (N / 2));
  681. if (!Wcq_alloc || !Wsq_alloc) {
  682. if (Wcq_alloc)
  683. luat_heap_free(Wcq_alloc);
  684. if (Wsq_alloc)
  685. luat_heap_free(Wsq_alloc);
  686. luat_heap_free(rq);
  687. luat_heap_free(iq);
  688. if (r_free)
  689. luat_heap_free(r);
  690. if (im_free)
  691. luat_heap_free(im);
  692. if (wc_free)
  693. luat_heap_free(Wc);
  694. if (ws_free)
  695. luat_heap_free(Ws);
  696. return luaL_error(L, "内存不足");
  697. }
  698. luat_fft_generate_twiddles_q15(Wcq_alloc, Wsq_alloc, N);
  699. Wcq = Wcq_alloc;
  700. Wsq = Wsq_alloc;
  701. }
  702. int scale_exp = 0;
  703. int rc = luat_fft_inplace_q15(rq, iq, N, 1, Wcq, Wsq, 0, &scale_exp); // inverse=1
  704. if (Wcq_alloc)
  705. luat_heap_free(Wcq_alloc);
  706. if (Wsq_alloc)
  707. luat_heap_free(Wsq_alloc);
  708. if (rc != 0) {
  709. luat_heap_free(rq);
  710. luat_heap_free(iq);
  711. if (r_free)
  712. luat_heap_free(r);
  713. if (im_free)
  714. luat_heap_free(im);
  715. if (wc_free)
  716. luat_heap_free(Wc);
  717. if (ws_free)
  718. luat_heap_free(Ws);
  719. return luaL_error(L, "q15 内核执行失败");
  720. }
  721. for (int i = 0; i < N; i++)
  722. ((int16_t*)zb_real->addr)[i] = rq[i];
  723. if (zb_imag && (int)zb_imag->len >= N * 2) {
  724. for (int i = 0; i < N; i++)
  725. ((int16_t*)zb_imag->addr)[i] = iq[i];
  726. }
  727. luat_heap_free(rq);
  728. luat_heap_free(iq);
  729. if (r_free)
  730. luat_heap_free(r);
  731. if (im_free)
  732. luat_heap_free(im);
  733. if (wc_free)
  734. luat_heap_free(Wc);
  735. if (ws_free)
  736. luat_heap_free(Ws);
  737. return 0;
  738. }
  739. luat_ifft_run_inplace(r, im, N, Wc, Ws);
  740. if (!luaL_testudata(L, 1, LUAT_ZBUFF_TYPE)) {
  741. lua_settop(L, 2);
  742. for (int i = 0; i < N; i++) {
  743. lua_pushnumber(L, r[i]);
  744. lua_rawseti(L, 1, i + 1);
  745. }
  746. }
  747. if (!luaL_testudata(L, 2, LUAT_ZBUFF_TYPE)) {
  748. for (int i = 0; i < N; i++) {
  749. lua_pushnumber(L, im[i]);
  750. lua_rawseti(L, 2, i + 1);
  751. }
  752. }
  753. if (r_free)
  754. luat_heap_free(r);
  755. if (im_free)
  756. luat_heap_free(im);
  757. if (wc_free)
  758. luat_heap_free(Wc);
  759. if (ws_free)
  760. luat_heap_free(Ws);
  761. return 0;
  762. }
  763. /*
  764. 频域积分(1/(jω))
  765. @api fft.fft_integral(real, imag, n, df)
  766. @param real 实部(float32,Lua 数组或 zbuff)
  767. @param imag 虚部(float32,Lua 数组或 zbuff)
  768. @int n 点数,2 的幂
  769. @number df 频率分辨率(fs/n)
  770. @return nil 原地修改 real/imag(DC 置 0)
  771. @usage
  772. -- 先完成 FFT 得到谱 (real, imag),再调用积分:
  773. fft.fft_integral(real, imag, N, fs/N)
  774. */
  775. static int l_fft_integral(lua_State* L)
  776. {
  777. int n = luaL_checkinteger(L, 3);
  778. float df = (float)luaL_checknumber(L, 4);
  779. if (n <= 1 || (n & (n - 1)))
  780. return luaL_error(L, "n must be power of 2");
  781. float *r = NULL, *im = NULL;
  782. int r_free = 0, im_free = 0;
  783. luat_zbuff_t* zb = NULL;
  784. zb = (luat_zbuff_t*)luaL_testudata(L, 1, LUAT_ZBUFF_TYPE);
  785. if (zb) {
  786. r = (float*)zb->addr;
  787. } else {
  788. r = luat_heap_malloc(sizeof(float) * n);
  789. r_free = 1;
  790. if (!r)
  791. return luaL_error(L, "no memory");
  792. if (read_lua_array_float(L, 1, r, n)) {
  793. if (r_free)
  794. luat_heap_free(r);
  795. return luaL_error(L, "real must be number array or zbuff");
  796. }
  797. }
  798. zb = (luat_zbuff_t*)luaL_testudata(L, 2, LUAT_ZBUFF_TYPE);
  799. if (zb) {
  800. im = (float*)zb->addr;
  801. } else {
  802. im = luat_heap_malloc(sizeof(float) * n);
  803. im_free = 1;
  804. if (!im) {
  805. if (r_free)
  806. luat_heap_free(r);
  807. return luaL_error(L, "no memory");
  808. }
  809. if (read_lua_array_float(L, 2, im, n)) {
  810. if (r_free)
  811. luat_heap_free(r);
  812. if (im_free)
  813. luat_heap_free(im);
  814. return luaL_error(L, "imag must be number array or zbuff");
  815. }
  816. }
  817. luat_fft_integral_inplace(r, im, n, df);
  818. if (!luaL_testudata(L, 1, LUAT_ZBUFF_TYPE)) {
  819. lua_settop(L, 2);
  820. for (int i = 0; i < n; i++) {
  821. lua_pushnumber(L, r[i]);
  822. lua_rawseti(L, 1, i + 1);
  823. }
  824. }
  825. if (!luaL_testudata(L, 2, LUAT_ZBUFF_TYPE)) {
  826. for (int i = 0; i < n; i++) {
  827. lua_pushnumber(L, im[i]);
  828. lua_rawseti(L, 2, i + 1);
  829. }
  830. }
  831. if (r_free)
  832. luat_heap_free(r);
  833. if (im_free)
  834. luat_heap_free(im);
  835. return 0;
  836. }
  837. static const rotable_Reg_t reg_fft[] = {
  838. { "generate_twiddles", ROREG_FUNC(l_fft_generate_twiddles) },
  839. { "generate_twiddles_q15_to_zbuff", ROREG_FUNC(l_fft_generate_twiddles_q15_to_zbuff) },
  840. { "run", ROREG_FUNC(l_fft_run) },
  841. { "ifft", ROREG_FUNC(l_fft_ifft) },
  842. { "fft_integral", ROREG_FUNC(l_fft_integral) },
  843. { NULL, ROREG_INT(0) }
  844. };
  845. LUAMOD_API int luaopen_fft(lua_State* L)
  846. {
  847. luat_newlib2(L, reg_fft);
  848. return 1;
  849. }