CustomOpBuilder通过AsdSipFFTOpRunner接入ASDSIP FFT算子

概述

Ascend Sip Boost (ASDSIP) 算子加速库是基于华为Ascend AI处理器，专门为信号处理领域而设计的算子库。

当用户需要使用ASDSIP加速库的FFT类算子，而MindSpore未提供相应算子接口时，用户可以使用自定义算子的方法快速接入使用。

在基于CustomOpBuilder的自定义算子中，MindSpore提供了 PyboostRunner 方便用户在动态图接入自定义算子。现在针对ASDSIP算子，MindSpore又额外提供了一套AsdSipFFTOpRunner用于把ASDSIP FFT算子的调用流程和动态图多级流水封装到一起。

用户基于 AsdSipFFTOpRunner类对接ASDSIP FFT算子时，仅需要提供Param，并调用Init接口初始化（即构造Operation），再调用Run接口即可执行ASDSIP算子。还可以直接调用 RunAsdSipFFTOp函数一键执行（函数内包含了Init和Run接口的调用）。

本指南以一个FftC2C为例，展示ASDSIP FFT算子的接入流程。完整代码请参阅代码仓库。

安装ASDSIP加速库

点这里查看安装教程

安装成功之后，需要激活ASDSIP加速库的环境变量：

source /usr/local/Ascend/nnal/asdsip/set_env.sh &> /dev/null

FftC2C算子接入

这里使用ms::pynative::RunAsdSipFFTOp接入算子，并通过ms::pynative::PyboostRunner::Call调用函数接口：

#include "ms_extension/api.h"

ms::Tensor InferFFTForward(const ms::Tensor &input) {
  ShapeVector out_tensor_shape(input.shape());
  return ms::Tensor(input.data_type(), out_tensor_shape);
}

ms::Tensor npu_fft(const ms::Tensor &input, int64_t n, int64_t batch_size) {
  ms::pynative::FFTParam param;
  param.fftXSize = n;
  param.fftYSize = 0;
  param.fftType = ms::pynative::asdFftType::ASCEND_FFT_C2C;
  param.direction = ms::pynative::asdFftDirection::ASCEND_FFT_FORWARD;
  param.batchSize = batch_size;
  param.dimType = ms::pynative::asdFft1dDimType::ASCEND_FFT_HORIZONTAL;
  auto output = InferFFTForward(input);
  ms::pynative::RunAsdSipFFTOp("asdFftExecC2C", param, input, output);
  return output;
}

auto pyboost_npu_fft(const ms::Tensor &input, int64_t n, int64_t batch_size) {
  return ms::pynative::PyboostRunner::Call<1>(npu_fft, input, n, batch_size);
}

PYBIND11_MODULE(MS_EXTENSION_NAME, m) {
  m.def("fft", &pyboost_npu_fft, "FFT C2C", pybind11::arg("input"), pybind11::arg("n"),
        pybind11::arg("batch_size"));
}

1. 推导算子的输出信息

ms::Tensor InferFFT1DForward(const ms::Tensor &input) {
  ShapeVector out_tensor_shape(input.shape());
  return ms::Tensor(input.data_type(), out_tensor_shape);
}

对于FftC2C算子，输出张量的数据类型和输入的一样。推导出输出形状之后，通过ms::Tensor构造函数构造一个空的张量。

2. 创建并设置算子属性结构体

ms::pynative::FFTParam param;
param.fftXSize = n;
param.fftYSize = 0;
param.fftType = ms::pynative::asdFftType::ASCEND_FFT_C2C;
param.direction = ms::pynative::asdFftDirection::ASCEND_FFT_FORWARD;
param.batchSize = batch_size;
param.dimType = ms::pynative::asdFft1dDimType::ASCEND_FFT_HORIZONTAL;

3. 调用RunAtbOp接口执行算子

ms::pynative::PyboostRunner::Call<1>(npu_fft, input, n, batch_size);

这是一个模板接口，其等效于：

auto runner = std::make_shared<AsdSipFFTOpRunner>("FFTExecC2C");
runner->Init(fft_param);
runner->Run({input}, {output});

传入算子名、属性、输入张量、输出张量几个信息，即可调用相应的ASDSIP算子。此接口支持了动态图的多级流水执行流程。

4. 通过pybind11将C++函数绑定一个Python函数

auto pyboost_npu_fft(const ms::Tensor &input, int64_t n, int64_t batch_size) {
  return ms::pynative::PyboostRunner::Call<1>(npu_fft, input, n, batch_size);
}

PYBIND11_MODULE(MS_EXTENSION_NAME, m) {
  m.def("fft", &pyboost_npu_fft, "FFT C2C", pybind11::arg("input"), pybind11::arg("n"),
        pybind11::arg("batch_size"));
}

5. 使用CustomOpBuilder编译自定义算子

将上述C++代码保存成文件asdsip_fftc2c.cpp，然后使用Python接口CustomOpBuilder编译。

input_np = np.random.rand(2, 16)
real_np = input_np.astype(np.float32)
imag_np = input_np.astype(np.float32)
complex_np = real_np + 1j * imag_np
my_ops = CustomOpBuilder("asdsip_fftc2c", "jit_test_files/asdsip_fftc2c.cpp", enable_asdsip=True).load()
output_tensor = my_ops.fft(input_tensor, 16, 2)
print(output_tensor)

这里向CustomOpBuilder传入了enable_asdsip=True的参数，MindSpore会自动添加与ASDSIP加速库有关的编译和链接选项。用户需保证正确执行了ASDSIP库的set_env.sh脚本，环境中已配置ASDSIP_HOME_PATH环境变量。