Closed shouhoo closed 5 months ago
[TOC]
二维FFT（Fast Fourier Transform）由一维FFT推广而来，近几十年来常被运用在图像处理上。DFT能够将二维复数序列{$${x_{n_1,n2}}$$}转化为二维复数序列{$${Y{n_1,n_2}}$$}，计算公式为：
可以将二维FFT看作是在输入的两个维度分别做一维FFT，因此其计算方式与一维FFT本质相同，这里不再对FFT计算细节进行赘述。
目前利用一维FFT以及转置算子实现二维FFT，如图1所示，其步骤如下，以[batch=1,n0, n1]为例：
1) 调用FFT 1D 算子，其参数为batch=n0, n = [n1] 2) 调用转置算子，数据类型为复数浮点数，将[n0, n1]转置为[n1, n0] 3) 调用FFT 1D 算子，其参数为batch=n1, n = [n0] 4) 调用转置算子，数据类型为复数浮点数，将[n1, n0]转置为[n0, n1]
图1 FFT2D 简单实现
当前与FFT 2D重点规模相关的FFT1D的实现性能如下：
表1 FFT1D 二的幂重点规模性能
表2 FF1D 非二的幂重点规模性能
测试[1, 2048, 6000]、[1, 2048, 14000]等2D规模，并与v100上cufft的性能对比：1)和3)两个FFT1D算子所占时间约为cufft的10倍；2)和4)两个转置算子所占时间约为cufft的40倍。
表3 FFT2D重点规模性能
表4 FFT2D重点规模性能(直接去除转置算子，运算结果错误)
FFT1D中的蝶形网络采用的是片下网络（GDRAM）-片上网络（NRAM）双层网络，通过将FFT规模分解为尽可能少的stage，每个stage中的蝶形加载至NRAM中进行处理。在NRAM中，即片上网络中，依然可以利用FFT将蝶形进一步分解，从而优化计算量。
FFT1D中的软流水优化仅优化了片上网络的性能，用片上网络的计算隐藏访存，而不同batch之间是几乎完全串行的。但对于2048和6000规模，由于整个FFT 1D的片下网络能够完全放到片上，即片下网络的stage为1，故只需加载一次片上网络并进行处理，故软流水优化无效。
对于简单版本实现来说，转置的耗时占比高达80%，这是不可接受的，因此应当进一步优化实现，直接去掉转置，由下一部分详细介绍。
想要去除两个转置算子，可以通过列主序FFT 1D实现，如图2所示，其FFT 2D计算步骤如下：
1) 调用FFT 1D 算子，其参数为batch=n0, n = [n1] 2) 调用列主序FFT 1D 算子，其参数为batch=n1, n = [n0]
在已实现的FFT 1D算子中，从片下网络加载大基至片上的方式是利用2D __memcpy进行gather-load操作，如图1所示。当片上网络规模未达到NRAM的容量，我们会同时加载同一个FFT中连续的多个大基，从而提升gather-load的效率。
而对于列主序版本，如果要加载连续的多个大基至片上，则需要加载不同batch的FFT中相同位置的大基，如图二所示。与行主序版本的gather-load和scatter-store相比，stride更大，性能可能明显下降，因此将采用分块与重排等方式提升数据的局部性。
图2 FFT2D 列主序版本实现
由于对于较小的FFT 1D (n≤6000)，增加针对不同batch之间计算隐藏访存的软流水。预期能够对重点规模中2048和6000相关的FFT2D性能进行提升。
FFT2D:实现对输入规模为[n0, n1]的二维复数序列做二维傅里叶变换。
IFFT2D: 实现对输入规模为[n0, n1]的二维复数序列做二维逆傅里叶变换。
与FFT1D接口相同
typedef struct mluopFFTStruct *mluopFFTPlan_t;
mluopStatus_t mluopCreateFFTPlan(mluopFFTPlan_t *fft_plan);
mluopStatus_t mluopMakeFFTPlanMany( mluopHandle_t handle, mluopFFTPlan_t fft_plan, const mluopTensorDescriptor_t input_desc, const mluopTensorDescriptor_t output_desc, const int rank, const int n[], size_t *reservespace_size, size_t *workspace_size);
mluopStatus_t mluopSetFFTReserveArea( mluopHandle_t handle, mluopFFTPlan_t fft_plan, void *reservespace);
mluopStatus_t mluopExecFFT(mluopHandle_t handle, const mluopFFTPlan_t fft_plan, const void *input, const float scale_factor, void *workspace, void *output, int direction);
mluopStatus_t mluopDestroyFFTPlan(mluopFFTPlan_t fft_plan);
在该实现方案中，我们首先简要描述了FFT2D的算法及其与FFT1D之间的关联，其次分析了调用FFT1D行主序算子与转置算子实现的简单版本，认为主要性能损失在于转置算子的调用以及部分规模软流水优化的实现，并根据该分析提出了相应的优化方法，即通过实现列主序FFT1D算子去除转置算子，并针对部分规模1D FFT进行batch间的软流水优化。最后给出需求分析与接口设计。
MLU 2D C2C FFT 方案
[TOC]
1 FFT 2D算法介绍
二维FFT（Fast Fourier Transform）由一维FFT推广而来，近几十年来常被运用在图像处理上。DFT能够将二维复数序列{$${x_{n_1,n2}}$$}转化为二维复数序列{$${Y{n_1,n_2}}$$}，计算公式为：
可以将二维FFT看作是在输入的两个维度分别做一维FFT，因此其计算方式与一维FFT本质相同，这里不再对FFT计算细节进行赘述。
2 MLU实现性能分析
2.1 简单实现
目前利用一维FFT以及转置算子实现二维FFT，如图1所示，其步骤如下，以[batch=1,n0, n1]为例：
1) 调用FFT 1D 算子，其参数为batch=n0, n = [n1] 2) 调用转置算子，数据类型为复数浮点数，将[n0, n1]转置为[n1, n0] 3) 调用FFT 1D 算子，其参数为batch=n1, n = [n0] 4) 调用转置算子，数据类型为复数浮点数，将[n1, n0]转置为[n0, n1]
图1 FFT2D 简单实现
2.2 性能测试
当前与FFT 2D重点规模相关的FFT1D的实现性能如下：
表1 FFT1D 二的幂重点规模性能
表2 FF1D 非二的幂重点规模性能
测试[1, 2048, 6000]、[1, 2048, 14000]等2D规模，并与v100上cufft的性能对比：1)和3)两个FFT1D算子所占时间约为cufft的10倍；2)和4)两个转置算子所占时间约为cufft的40倍。
表3 FFT2D重点规模性能
表4 FFT2D重点规模性能(直接去除转置算子，运算结果错误)
2.3 性能分析
FFT1D中的蝶形网络采用的是片下网络（GDRAM）-片上网络（NRAM）双层网络，通过将FFT规模分解为尽可能少的stage，每个stage中的蝶形加载至NRAM中进行处理。在NRAM中，即片上网络中，依然可以利用FFT将蝶形进一步分解，从而优化计算量。
FFT1D中的软流水优化仅优化了片上网络的性能，用片上网络的计算隐藏访存，而不同batch之间是几乎完全串行的。但对于2048和6000规模，由于整个FFT 1D的片下网络能够完全放到片上，即片下网络的stage为1，故只需加载一次片上网络并进行处理，故软流水优化无效。
对于简单版本实现来说，转置的耗时占比高达80%，这是不可接受的，因此应当进一步优化实现，直接去掉转置，由下一部分详细介绍。
3 MLU上的FFT 2D优化实现
3.1 列主序FFT 1D实现
想要去除两个转置算子，可以通过列主序FFT 1D实现，如图2所示，其FFT 2D计算步骤如下：
1) 调用FFT 1D 算子，其参数为batch=n0, n = [n1] 2) 调用列主序FFT 1D 算子，其参数为batch=n1, n = [n0]
在已实现的FFT 1D算子中，从片下网络加载大基至片上的方式是利用2D __memcpy进行gather-load操作，如图1所示。当片上网络规模未达到NRAM的容量，我们会同时加载同一个FFT中连续的多个大基，从而提升gather-load的效率。
而对于列主序版本，如果要加载连续的多个大基至片上，则需要加载不同batch的FFT中相同位置的大基，如图二所示。与行主序版本的gather-load和scatter-store相比，stride更大，性能可能明显下降，因此将采用分块与重排等方式提升数据的局部性。
图2 FFT2D 列主序版本实现
3.2 软流水优化
由于对于较小的FFT 1D (n≤6000)，增加针对不同batch之间计算隐藏访存的软流水。预期能够对重点规模中2048和6000相关的FFT2D性能进行提升。
4 MLU层需求分析
4.1 算子需求分析
4.2 算子功能和应用场景描述
FFT2D:实现对输入规模为[n0, n1]的二维复数序列做二维傅里叶变换。
IFFT2D: 实现对输入规模为[n0, n1]的二维复数序列做二维逆傅里叶变换。
4.3 算子输入输出参数要求
5算子接口设计
与FFT1D接口相同
6 总结
在该实现方案中，我们首先简要描述了FFT2D的算法及其与FFT1D之间的关联，其次分析了调用FFT1D行主序算子与转置算子实现的简单版本，认为主要性能损失在于转置算子的调用以及部分规模软流水优化的实现，并根据该分析提出了相应的优化方法，即通过实现列主序FFT1D算子去除转置算子，并针对部分规模1D FFT进行batch间的软流水优化。最后给出需求分析与接口设计。