CANN ops-math 中的 aclStft 算子：NPU 上滑动窗口傅里叶变换（STFT）的两段式接口调用与实现解析

CANN ops-math 中的 aclStft 算子：NPU 上滑动窗口傅里叶变换（STFT）的两段式接口调用与实现解析

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math

本文以 math/stft/docs/aclStft.md 为核心，系统讲解 CANN ops-math 仓库中 STFT 算子的 aclnn 接口（aclStft）的功能定义、两段式调用流程、全部入参约束与返回码语义，并结合 op_host/op_api/acl_stft.cpp 与 op_kernel/stft.cpp 等源码，剖析其"DFT 矩阵 + Matmul"的实现原理、AiCore/AiCpu 双路径调度与 DFT 矩阵缓存机制。读完本文，你将能够独立完成 aclStft 的 workspace 计算、执行器创建、流式下发与结果回收，并理解其输出 shape 推导、数据类型对应关系与性能相关约束。

一、功能说明与数学定义

STFT（Short-Time Fourier Transform，短时傅里叶变换）将长时序信号切分为若干滑动窗口，再对每个窗口分别做傅里叶变换，得到"时间-频率"二维表示。aclStft接口计算输入在滑动窗口内的傅里叶变换，与 PyTorch 的torch.stft对齐（可参考 math/stft/op_graph/stft_proto.h 中 "Compatible with pytorch STFT operator" 的注释）。

当normalized=False时：

$$ X[w,m]=\sum_{k=0}^{winLength-1}window[k]*self[m*hopLength+k]exp(-j\frac{2{\pi}wk}{nFft}) $$

当normalized=True时：

$$ X[w,m]=\frac{1}{\sqrt{nFft}}(\sum_{k=0}^{winLength-1}window[k]*self[m*hopLength+k]exp(-j\frac{2{\pi}wk}{nFft})) $$

其中：

• $w$：FFT 的频点（frequency bin）。
• $m$：滑动窗口的 index。
• $self$：1 维或 2 维 Tensor。1 维时表示一条时序采样序列；2 维时表示多条时序采样序列（batch 维）。
• $hopLength$：滑动窗口的步进间隔（hop）。
• $window$：1 维 Tensor，即 STFT 的窗函数（如 hann_window），长度为 $winLength$。
• $exp(-j*\frac{2{\pi}wk}{nFft})$：旋转因子（twiddle factor）。

从源码视角看，该公式在 NPU 上并非逐点循环实现，而是被等价转化为"DFT 矩阵 × 分帧矩阵"的矩阵乘法问题：aclStftGetWorkspaceSize内部生成旋转因子矩阵（DFT matrix）$W[k,n]=exp(-j*2\pi k n / nFft)$，配合窗函数与输入分帧结果，通过 AiCore 上的 Matmul 完成频域计算（详见后文"源码级实现"一节）。

二、产品支持情况

aclStft在不同 NPU 平台上的支持情况如下：

该表与算子级 README（math/stft/README.md）完全一致，且在源码中有直接印证：acl_stft.cpp中的CheckPlatform()仅放行ASCEND910B、ASCEND910_93（对应 A3）、ASCEND950三类 SoC，其他平台直接返回ACLNN_ERR_PARAM_INVALID并打印 "STFT is not supported on this platform"；算子定义 op_host/stft_def.cpp 中也只为ascend910b、ascend910_93、ascend950三个平台注册了 AICore 配置。

三、两段式接口与函数原型

aclStft属于 CANN 的两段式接口：必须先调用aclStftGetWorkspaceSize获取计算所需 workspace 大小以及包含算子计算流程的执行器（executor），再调用aclStft执行计算。

第一段接口原型：

aclnnStatus aclStftGetWorkspaceSize( const aclTensor *self, const aclTensor *windowOptional, aclTensor *out, int64_t nFft, int64_t hopLength, int64_t winLength, bool normalized, bool onesided, bool returnComplex, uint64_t *workspaceSize, aclOpExecutor **executor)

第二段接口原型：

aclnnStatus aclStft( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

对应的头文件声明位于 math/stft/op_host/op_api/acl_stft.h，其 doxygen 注释中还给出了内部计算图：self → l0op::Stft → l0op::ViewCopy → out，即算子计算结果最终通过 ViewCopy 落到用户提供的outTensor 上。

四、aclStftGetWorkspaceSize 参数详解

第一段接口承担了全部入参校验、计算图构建与 workspace 估算任务。完整参数说明如下：

4.1 输出 shape 的推导规则

输出维度由输入维度、onesided与returnComplex共同决定，源码GetOutputShape()（见 math/stft/op_host/op_api/acl_stft.cpp）中的推导逻辑与文档一致：

• 滑动窗口个数：frames = (L - nFft) / hopLength + 1
• 频率维：n = onesided ? nFft / 2 + 1 : nFft
• returnComplex=True：1 维输入 → [n, frames]；2 维输入 → [B, n, frames]
• returnComplex=False：1 维输入 → [n, frames, 2]；2 维输入 → [B, n, frames, 2]

也就是说，out的实际 shape 完全由self的 shape 与三个布尔属性推导而来，调用方不能随意指定，CheckShape()中会通过OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE严格校验。

4.2 与算子级 README 参数表的对应关系

在算子级文档 math/stft/README.md 中，同样的参数以x / window / n_fft / hop_length / win_length / normalized / onesided / return_complex / y命名，并补充了图模式下属性的默认值，可供 aclnn 接口使用参考：

在算子 IR 定义 math/stft/op_graph/stft_proto.h 中，这些默认值同样被登记：hop_length=0、win_length=0（0 表示"未指定，按默认规则取 floor(n_fft/4) 与 n_fft"）、normalized=false、onesided=true、return_complex=true，n_fft为必填属性。需要留意的是，虽然仓库中注册了 STFT 的算子 IR，但 math/stft/op_host/stft_def.cpp 中StftCheckSupport当前返回 "Stft don't support graph mode"，即从当前源码看，图模式调用并未放开，aclnn 两段式接口是推荐的调用方式（README 的"调用说明"一节也将 aclnn 接口列为第一种调用方式）。

五、返回码与错误场景

aclStftGetWorkspaceSize返回aclnnStatus状态码，具体取值可参见 aclnn 返回码说明。第一段接口完成入参校验，以下场景会报错：

这些校验在源码CheckParams()及其子函数中逐条实现，顺序为：空指针检查（CheckNotNull）→ 数据类型检查（CheckDtypeValid，含self/window类型一致性）→ 数据格式检查（CheckFormat，要求 ND）→ shape 与属性检查（CheckShape）。其中CheckShape还会校验：

• nFft > 0且len >= nFft；
• hopLength > 0；
• 0 < winLength <= nFft；
• 当winLength != nFft时，window 的 shape 第 0 维必须等于winLength；
• 输入为复数时onesided必须为 False。

六、aclStft 参数说明

第二段接口仅负责执行，参数如下：

其实现（aclStft函数体）非常精简，仅包含L2_DFX_PHASE_2性能打点与CommonOpExecutorRun(workspace, workspaceSize, executor, stream)的固定调用，真正的工作都在第一段接口中完成。

七、约束说明

• 与 PyTorch 输入的差异：PyTorch 接口的输入 self 为原始输入；aclStftGetWorkspaceSize的入参 self 是原始输入经过前端 PyTorch 补 pad 后得到的结果。这一点意味着用户在 aclnn 侧直接调用时，需要自行保证输入长度满足窗口切分要求。

• 大 shape 超时风险：当输入 self 的 shape 为 [B, L]，且下式计算的结果较大时，当前接口的计算可能超时：

  $$ B * ((L - nFft) / hopLength + 1) * nFft $$

  从源码看，该公式对应分帧后的总数据规模：AIC 侧 Matmul 的 M 维 × 帧数 × batch，规模过大时核内循环（frame split + matmul）耗时显著上升。

• 约束：nFft <= L；winLength <= nFft。

• normalized 语义：当 normalized=True 时，输出等价于对原始结果乘以 $\frac{1}{\sqrt{N}}$：

  $$ STFT(w,m)=\frac{1}{\sqrt{N}}X[w,m] $$

• 数据类型对应关系：self、windowOptional、returnComplex、out 之间的对应关系如下表：

  self	windowOptional	returnComplex	out
  FLOAT32	FLOAT32	True	COMPLEX64
  DOUBLE	DOUBLE	True	COMPLEX128
  COMPLEX64	COMPLEX64	True	COMPLEX64
  COMPLEX128	COMPLEX128	True	COMPLEX128
  FLOAT32	FLOAT32	False	FLOAT32
  DOUBLE	DOUBLE	False	DOUBLE
  COMPLEX64	COMPLEX64	False	FLOAT32
  COMPLEX128	COMPLEX128	False	DOUBLE

  即：returnComplex=True 时，实数输入输出升级为复数类型；returnComplex=False 时，复数输入退化为实/虚部各占一维的实数输出（out 多出最后一维 size=2）。

• 确定性计算：aclStft默认是确定性实现（同一输入多次执行结果一致）。

八、调用示例

示例代码如下，来源于 math/stft/examples/test_aclnn_stft.cpp（与文档示例等价，且采用 RAII 智能指针管理资源）。具体编译和执行过程可参考编译与运行样例。

#include <iostream> #include <vector> #include "acl/acl.h" #include "aclnnop/acl_stft.h" #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vector<int64_t>& shape) { int64_t shapeSize = 1; for (auto i : shape) { shapeSize *= i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法，资源初始化 auto ret = aclInit(nullptr); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret); ret = aclrtSetDevice(deviceId); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret); ret = aclrtCreateStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret); return 0; } template <typename T> int CreateAclTensor( const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size = GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret); // 计算连续tensor的strides std::vector<int64_t> strides(shape.size(), 1); for (int64_t i = shape.size() - 2; i >= 0; i--) { strides[i] = shape[i + 1] * strides[i + 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor = aclCreateTensor( shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1.（固定写法）device/stream初始化，参考acl对外接口列表 // 根据自己的实际device填写deviceId int32_t deviceId = 0; aclrtStream stream; auto ret = Init(deviceId, &stream); // check根据自己的需要处理 CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret); // 2.构造输入与输出，需要根据API的接口自定义构造 std::vector<int64_t> selfShape = {5}; std::vector<int64_t> windowShape = {4}; std::vector<int64_t> outShape = {3, 1, 2}; void* selfDeviceAddr = nullptr; void* windowDeviceAddr = nullptr; void* outDeviceAddr = nullptr; aclTensor* self = nullptr; aclTensor* window = nullptr; aclTensor* out = nullptr; std::vector<float> selfHostData = {1, 6, 8, 5, 7}; std::vector<float> windowHostData = {1, 1, 1, 1}; std::vector<float> outHostData = {0, 0, 0, 0, 0, 0}; // 创建self aclTensor ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建window aclTensor ret = CreateAclTensor(windowHostData, windowShape, &windowDeviceAddr, aclDataType::ACL_FLOAT, &window); CHECK_RET(ret == ACL_SUCCESS, return ret); // 创建out aclTensor ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out); CHECK_RET(ret == ACL_SUCCESS, return ret); int n_fft = 4; int hop_length = 2; int win_length = 4; bool normalized = false; bool onesided = true; bool returnComplex = false; // 3.调用CANN算子库API，需要修改为具体的API名称 uint64_t workspaceSize = 0; aclOpExecutor* executor; // 调用aclStft第一段接口 ret = aclStftGetWorkspaceSize( self, window, out, n_fft, hop_length, win_length, normalized, onesided, returnComplex, &workspaceSize, &executor); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclStftGetWorkspaceSize failed. ERROR: %d\n", ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr = nullptr; if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret); } // 调用aclStft第二段接口 ret = aclStft(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclStft failed. ERROR: %d\n", ret); return ret); // 4.（固定写法）同步等待任务执行结束 ret = aclrtSynchronizeStream(stream); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret); // 5.获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改 auto size = GetShapeSize(outShape); std::vector<float> resultData(size, 0); ret = aclrtMemcpy( resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret); for (int64_t i = 0; i < size; i++) { LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]); } // 6.释放aclTensor，需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(window); aclDestroyTensor(out); // 7.释放device资源，需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(windowDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize > 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }

8.1 示例参数与输出 shape 的对应验证

示例中selfShape = {5}（L=5）、n_fft = 4、hop_length = 2、win_length = 4、onesided = true、returnComplex = false，则：

• 频率维 N = n_fft // 2 + 1 = 3；
• 窗口数 T = (5 - 4) // 2 + 1 = 1；
• returnComplex=False，故 out 的 shape 为 [N, T, 2] = [3, 1, 2]，与代码中的outShape = {3, 1, 2}完全吻合。

调用流程可归纳为 7 步固定套路：① acl 环境初始化与 stream 创建；② 构造 self/window/out 三个 aclTensor（host→device 拷贝 +aclCreateTensor）；③ 第一段接口求 workspaceSize 与 executor，并按需aclrtMallocworkspace；④ 第二段接口aclStft异步下发；⑤aclrtSynchronizeStream同步等待；⑥aclrtMemcpy回拷结果并打印；⑦ 依次释放 tensor、device 内存、workspace、stream 并aclFinalize。

九、源码级实现剖析

9.1 第一段接口的整体流程

aclStftGetWorkspaceSize的执行顺序（对应 math/stft/op_host/op_api/acl_stft.cpp）：

1. L2_DFX_PHASE_1记录入参（self、windowOptional、nFft、hopLength、winLength、normalized、onesided、returnComplex）与输出 out，用于问题定位；
2. CREATE_EXECUTOR()创建 OpExecutor；
3. CheckPlatform()平台校验（仅 910B/910_93/950）；
4. CheckParams()完成空指针、dtype、format、shape 四类校验；
5. 空 Tensor 处理：self为空时直接返回，workspaceSize 置 0；
6. 通过l0op::Contiguous将 self 转为连续 Tensor；
7. 路径选择：调用l0op::IsStftAiCoreSupported(...)判断走 AiCore 高性能路径还是 AiCpu 兜底路径；
8. 构建计算图（详见 9.2），最后通过uniqueExecutor->GetWorkspaceSize()汇总 workspace 大小并ReleaseTo(executor)返回。

9.2 AiCore 与 AiCpu 双路径

从源码日志分支可以清晰看到两条路径：

• AiCpu 路径（path=AiCpu）：直接以l0op::Stft构造算子计算，再接l0op::ViewCopy将结果搬运到 out。用于不支持 AiCore 加速的 shape/类型组合；
• AiCore 路径（path=AiCore）：以矩阵乘法为核心，流程为：
  1. 若winLength < nFft（准确说是小于按 block 对齐后的nFftAlign），调用GeneratePadWindow用 PadV3（conversion/pad_v3算子，constant 模式、pad value=0）把窗函数补齐到 nFft 长度；若调用方未传 window，则先用l0op::OnesLike生成全 1 窗；
  2. 调用GenerateDftMatrix生成旋转因子矩阵（K × nFft，K = onesided ? nFft/2+1 : nFft），实部/虚部按频率行交错排布为 [2, K, nFft_align]；
  3. 组合l0op::Mul（把窗函数乘入 DFT 矩阵）与l0op::Stft完成分帧矩阵与 DFT 矩阵的乘加，最终l0op::ViewCopy落盘 out。

9.3 DFT 矩阵缓存（DftMatrixCache）

由于旋转因子矩阵只由 K、nFft 与对齐方式决定，acl_stft.cpp实现了基于 LRU 的 DFT 矩阵缓存：

• 缓存键为{K, nFft, nfftAlignBytes, deviceId}四元组，注释明确指出该 4 字段键可完整覆盖 8 个入参组合的所有场景（hopLength/winLength/normalized/onesided/returnComplex 只通过nfftAlignBytes间接影响矩阵布局）；
• 缓存预算默认 8GB（DEFAULT_DFT_CACHE_MAX_MEMORY），注释说明"可覆盖 nFft 到 32768 的所有常见尺度"；
• 单矩阵超过预算 50% 时不缓存（避免大矩阵挤占小矩阵）；命中缓存时通过SetStorageAddr复用已有显存并AbandonCache；
• 矩阵构造通过八分圆（octant）对称法计算 cos/sin，避免大范围三角函数计算。

从工程视角看，这对语音/音频类高频场景（同一 nFft 反复调用）可显著减少重复建矩阵的开销。

9.4 Kernel 侧实现：分帧 + Matmul

Kernel 入口stft（见 math/stft/op_kernel/stft.cpp）依据 tiling key 分发到 4 种实现：

以StftND（math/stft/op_kernel/stft.h）为例，其核心思路是 AIV/AIC 协同：

• AIV（Vector 核）：SplitFrameNormal按 hop 步进把输入切分为blkFrame × nfft的帧矩阵写入 workspace（分帧窗口重叠拷贝）；随后把 Matmul 输出的实部/虚部通过Gather指令交错拼成复数排列写回输出；
• AIC（Cube 核）：以"分帧矩阵 × DFT 矩阵"为 Matmul 主体（M 为频率行/窗函数行，N 为帧数，K 为 nFft），通过Mmad/Fixpipe完成累加与落盘，采用 L1/A0/B0/C0 多级流水与双缓冲（baseM=96、baseK=80、N 基本块 96 等经验参数）提升吞吐。

9.5 Tiling 与算子定义

• Tiling 入口在 math/stft/op_host/stft_tiling.cpp，通过Math::OpTiling::TilingRegistry分派具体 tiling 实现，TilingPrepare4STFT在编译期采集核数（AIV/AIC）、UB 内存大小、系统 workspace 等平台信息；
• 算子定义在 math/stft/op_host/stft_def.cpp，输入为x（必填）、plan（必填，即 DFT 矩阵/plan）、window（可选），输出y，并开启了动态 rank、动态 shape、dynamic compile、precision reduce 等特性；
• 由于 aclnn 接口内部使用 l0op 拼装计算图（Contiguous/PadV3/Mul/Stft/ViewCopy），用户侧只需面对aclStft一个入口，复杂度被完全封装。

十、测试与验证

仓库为 aclStft 提供了完整的测试覆盖，可作为功能与正确性验证的参考：

• 算子级 ST 测试：math/stft/tests/st/aclStft/atk_aclStft.json 以torch.stft为基准（name: "torch.stft"，aclnn_name: "aclStft"），覆盖数十组用例：输入 shape 为 [16, 约 10 万～23 万] 的长序列、FLOAT32、nFft=400、hopLength=160、winLength=400、onesided=true、returnComplex=false，数值范围 [-10.0, 10.0]，精度标准为cv_fused_double_benchmark。这组参数恰与源码中NfftAlignBytes的特化分支（nFft=400、hop=160、onesided、非复数，走 32B block 对齐的 X1 性能路径）对应；
• Host 侧 UT：math/stft/tests/ut/op_host/op_api/test_aclnn_stft.cpp 验证 aclnn 接口调用；math/stft/tests/ut/op_host/test_stft_infershape.cpp 与 math/stft/tests/ut/op_host/test_stft_tiling.cpp 分别验证 shape 推导与 tiling 计算；
• Kernel 侧 UT：math/stft/tests/ut/op_kernel/test_stft.cpp 配合stft_tiling.h验证 kernel 计算正确性；
• 可直接编译运行的示例：math/stft/examples/test_aclnn_stft.cpp 即第八节代码的完整 RAII 版本，适合作为上板验证的最小可运行工程。

十一、相关参考

• 本文核心文档：math/stft/docs/aclStft.md
• 算子级说明与参数表：math/stft/README.md
• aclnn 对外头文件：math/stft/op_host/op_api/acl_stft.h
• 接口实现与 DFT 矩阵缓存：math/stft/op_host/op_api/acl_stft.cpp
• Kernel 实现：math/stft/op_kernel/stft.cpp、math/stft/op_kernel/stft.h
• 算子定义与 IR：math/stft/op_host/stft_def.cpp、math/stft/op_graph/stft_proto.h
• 通用概念：两段式接口、aclnn 返回码、编译与运行样例

【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库，实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math