muFFT 开发者指南

muFFT 是基于 MUSA 开发的离散傅里叶变换库，在 MTGPU 上经过深度优化，在深度学习、计算物理、分子动力学、量子化学、地震和医学成像等领域被广泛使用。

概述

什么是 muFFT

muFFT（MUSA Fast Fourier Transform，快速傅里叶变换库）提供基于 MUSA 的快速傅里叶变换接口，支持在 MTGPU 上执行一维、二维和三维 FFT，可用于信号处理、科学计算和深度学习等工作负载。

关键特性

• 支持 R2C、C2C、C2R 等常用 FFT 变换类型。
• 支持一维、二维、三维以及批处理 FFT。
• 提供 plan 创建、执行和参数配置接口，便于复用变换配置。
• 针对 MTGPU 架构优化，适合高吞吐频域计算场景。

支持的变换类型

变换类型	说明	精度
R2C	实数到复数 (Real to Complex)	单精度/双精度
C2C	复数到复数 (Complex to Complex)	单精度/双精度
C2R	复数到实数 (Complex to Real)	单精度/双精度

支持的维度

• 一维 FFT (1D)
• 二维 FFT (2D)
• 三维 FFT (3D)

批处理支持

支持批处理 FFT，可同时处理多个相同尺寸的变换。

---

快速开始

一维复数到复数变换 (1D C2C)

#include <musa_runtime.h>
#include <mufft.h>
#include <complex>
#include <iostream>
#include <vector>

int main() {
    std::cout << "muFFT 1D single-precision complex-to-complex transform\n";

    const int Nx = 8;
    std::vector<std::complex<float>> h_x(Nx);
    
    // 初始化输入数据
    for (size_t i = 0; i < Nx; i++) {
        h_x[i] = std::complex<float>(i, 0);
    }
    
    std::cout << "Input: ";
    for (size_t i = 0; i < Nx; i++) {
        std::cout << h_x[i] << " ";
    }
    std::cout << std::endl;

    // 创建 MUSA 设备对象并复制数据到设备
    size_t complex_bytes = sizeof(decltype(h_x)::value_type) * h_x.size();
    mufftComplex* d_x;
    musaMalloc(&d_x, complex_bytes);
    musaMemcpy(d_x, h_x.data(), complex_bytes, musaMemcpyHostToDevice);

    // 创建 plan
    mufftHandle plan;
    mufftPlan1d(&plan,           // plan handle
                Nx,              // 变换长度
                MUFFT_C2C,       // 变换类型 (MUFFT_C2C 单精度)
                1);              // 变换数量

    // 执行 plan (正向变换)
    mufftExecC2C(plan, d_x, d_x, MUFFT_FORWARD);

    // 复制结果回主机
    musaMemcpy(h_x.data(), d_x, complex_bytes, musaMemcpyDeviceToHost);
    
    std::cout << "Output: ";
    for (size_t i = 0; i < Nx; i++) {
        std::cout << h_x[i] << " ";
    }
    std::cout << std::endl;

    // 释放资源
    mufftDestroy(plan);
    musaFree(d_x);

    return 0;
}

编译和运行

# 使用 mcc 编译
mcc fft1d_example.cpp \
    -L/usr/local/musa/lib \
    -I/usr/local/musa/include \
    -lmufft -lmusart \
    -o fft1d_example

# 运行
./fft1d_example

---

API 参考

创建 Plan

一维 FFT Plan

mufftResult_t mufftPlan1d(mufftHandle* plan,
                          int nx,
                          mufftType_t type,
                          int batch);

参数:

参数	说明
plan	输出的 plan handle
nx	X 维度的尺寸
type	变换类型 (MUFFT_C2C, MUFFT_R2C, MUFFT_C2R)
batch	批处理数量

二维 FFT Plan

mufftResult_t mufftPlan2d(mufftHandle* plan,
                          int nx,
                          int ny,
                          mufftType_t type);

三维 FFT Plan

mufftResult_t mufftPlan3d(mufftHandle* plan,
                          int nx,
                          int ny,
                          int nz,
                          mufftType_t type);

执行变换

C2C 变换

mufftResult_t mufftExecC2C(mufftHandle plan,
                           mufftComplex* idata,
                           mufftComplex* odata,
                           int direction);

参数:

参数	说明
plan	FFT plan handle
idata	输入数据 (设备指针)
odata	输出数据 (设备指针)
direction	方向 (MUFFT_FORWARD 或 MUFFT_INVERSE)

R2C 变换

mufftResult_t mufftExecR2C(mufftHandle plan,
                           float* idata,
                           mufftComplex* odata);

C2R 变换

mufftResult_t mufftExecC2R(mufftHandle plan,
                           mufftComplex* idata,
                           float* odata);

设置参数

设置 Scale

mufftResult_t mufftSetScale(mufftHandle plan, float scale);

设置逆变换的缩放因子。

设置 Stream

mufftResult_t mufftSetStream(mufftHandle plan, musaStream_t stream);

绑定 FFT 操作到指定的 MUSA stream。

---

使用示例

后续示例的编译和链接方式与上文「编译和运行」中的示例相同，请参阅快速开始中的编译和运行说明。

二维 FFT

#include <musa_runtime.h>
#include <mufft.h>
#include <complex>
#include <vector>

int main() {
    const int Nx = 64;
    const int Ny = 64;
    
    std::vector<std::complex<float>> h_data(Nx * Ny);
    
    // 初始化数据
    for (int i = 0; i < Nx * Ny; i++) {
        h_data[i] = std::complex<float>(i % 10, 0);
    }
    
    // 分配设备内存
    size_t bytes = sizeof(std::complex<float>) * Nx * Ny;
    mufftComplex* d_data;
    musaMalloc(&d_data, bytes);
    musaMemcpy(d_data, h_data.data(), bytes, musaMemcpyHostToDevice);
    
    // 创建 2D FFT plan
    mufftHandle plan;
    mufftPlan2d(&plan, Nx, Ny, MUFFT_C2C);
    
    // 执行正向 FFT
    mufftExecC2C(plan, d_data, d_data, MUFFT_FORWARD);
    
    // 复制结果回主机
    musaMemcpy(h_data.data(), d_data, bytes, musaMemcpyDeviceToHost);
    
    // 清理
    mufftDestroy(plan);
    musaFree(d_data);
    
    return 0;
}

批处理 FFT

#include <musa_runtime.h>
#include <mufft.h>
#include <complex>
#include <vector>

int main() {
    const int Nx = 256;
    const int batch = 32;  // 批处理 32 个 FFT
    
    std::vector<std::complex<float>> h_data(Nx * batch);
    
    // 初始化数据
    for (int i = 0; i < Nx * batch; i++) {
        h_data[i] = std::complex<float>(i * 0.1f, 0);
    }
    
    // 分配设备内存
    size_t bytes = sizeof(std::complex<float>) * Nx * batch;
    mufftComplex* d_data;
    musaMalloc(&d_data, bytes);
    musaMemcpy(d_data, h_data.data(), bytes, musaMemcpyHostToDevice);
    
    // 创建批处理 1D FFT plan
    mufftHandle plan;
    mufftPlan1d(&plan, Nx, MUFFT_C2C, batch);
    
    // 执行批处理 FFT
    mufftExecC2C(plan, d_data, d_data, MUFFT_FORWARD);
    
    // 复制结果回主机
    musaMemcpy(h_data.data(), d_data, bytes, musaMemcpyDeviceToHost);
    
    // 清理
    mufftDestroy(plan);
    musaFree(d_data);
    
    return 0;
}

实数 FFT (R2C + C2R)

#include <musa_runtime.h>
#include <mufft.h>
#include <complex>
#include <vector>
#include <cmath>
#include <iostream>

int main() {
    const int Nx = 128;
    
    // 实数输入数据
    std::vector<float> h_real(Nx);
    std::vector<std::complex<float>> h_complex(Nx/2 + 1);  // R2C 输出尺寸为 Nx/2+1
    
    // 初始化实数数据 (正弦波)
    for (int i = 0; i < Nx; i++) {
        h_real[i] = std::sin(2.0f * 3.14159f * i / 16.0f);
    }
    
    // 分配设备内存
    float* d_real;
    mufftComplex* d_complex;
    
    size_t real_bytes = sizeof(float) * Nx;
    size_t complex_bytes = sizeof(std::complex<float>) * (Nx/2 + 1);
    
    musaMalloc(&d_real, real_bytes);
    musaMalloc(&d_complex, complex_bytes);
    
    // 复制实数数据到设备
    musaMemcpy(d_real, h_real.data(), real_bytes, musaMemcpyHostToDevice);
    
    // 创建 R2C plan
    mufftHandle plan_r2c;
    mufftPlan1d(&plan_r2c, Nx, MUFFT_R2C, 1);
    
    // 执行 R2C 变换
    mufftExecR2C(plan_r2c, d_real, d_complex);
    
    // 复制复数结果回主机
    musaMemcpy(h_complex.data(), d_complex, complex_bytes, musaMemcpyDeviceToHost);
    
    std::cout << "R2C Result (first 10 elements):\n";
    for (int i = 0; i < 10; i++) {
        std::cout << h_complex[i] << std::endl;
    }
    
    // C2R 逆变换
    // 首先复制数据到设备
    musaMemcpy(d_complex, h_complex.data(), complex_bytes, musaMemcpyHostToDevice);
    
    mufftHandle plan_c2r;
    mufftPlan1d(&plan_c2r, Nx, MUFFT_C2R, 1);
    
    // 执行 C2R 逆变换
    mufftExecC2R(plan_c2r, d_complex, d_real);
    
    // 复制实数结果回主机
    musaMemcpy(h_real.data(), d_real, real_bytes, musaMemcpyDeviceToHost);
    
    std::cout << "\nC2R Result (first 10 elements):\n";
    for (int i = 0; i < 10; i++) {
        std::cout << h_real[i] << std::endl;
    }
    
    // 清理
    mufftDestroy(plan_r2c);
    mufftDestroy(plan_c2r);
    musaFree(d_real);
    musaFree(d_complex);
    
    return 0;
}

---

性能优化建议

选择合适的尺寸

muFFT 在某些尺寸下性能更好：

• 2 的幂次：128, 256, 512, 1024...
• 小因子分解：3×2^n, 5×2^n

避免大质数尺寸。

批处理优于多次调用

// ❌ 性能较差：多次调用
for (int i = 0; i < batch; i++) {
    mufftExecC2C(plan, d_data + i*Nx, d_data + i*Nx, MUFFT_FORWARD);
}

// ✅ 性能较好：批处理 plan
mufftHandle batchPlan;
mufftPlan1d(&batchPlan, Nx, MUFFT_C2C, batch);
mufftExecC2C(batchPlan, d_data, d_data, MUFFT_FORWARD);

使用 Stream 并发

musaStream_t stream1, stream2;
musaStreamCreate(&stream1);
musaStreamCreate(&stream2);

mufftHandle plan1, plan2;
mufftPlan1d(&plan1, Nx, MUFFT_C2C, 1);
mufftPlan1d(&plan2, Nx, MUFFT_C2C, 1);

mufftSetStream(plan1, stream1);
mufftSetStream(plan2, stream2);

// 并发执行
mufftExecC2C(plan1, d_data1, d_out1, MUFFT_FORWARD);
mufftExecC2C(plan2, d_data2, d_out2, MUFFT_FORWARD);

---

错误处理

mufftResult_t result = mufftPlan1d(&plan, Nx, MUFFT_C2C, 1);

if (result != MUFFT_SUCCESS) {
    fprintf(stderr, "muFFT error: %d\n", result);
    
    switch (result) {
        case MUFFT_ALLOC_FAILED:
            fprintf(stderr, "Memory allocation failed\n");
            break;
        case MUFFT_INVALID_SIZE:
            fprintf(stderr, "Invalid FFT size\n");
            break;
        case MUFFT_INVALID_TYPE:
            fprintf(stderr, "Invalid FFT type\n");
            break;
        case MUFFT_EXEC_FAILED:
            fprintf(stderr, "Kernel execution failed\n");
            break;
        default:
            fprintf(stderr, "Unknown error\n");
    }
}

---

相关文档

• muFFT API 参考