自定义Ascend C算子，支持模型在华为NPU加速卡上部署运行

背景介绍

        近期，在对接客户需求时，对方需要在华为NPU上部署特定的深度学习模型，实现私有化交付。然而，目前我们现有的部署方案基本都是机遇Nvidia CUDA生态，尚未涉及华为Ascend方面。于是，我们决定基于Ascend生态，重新适配深度学习模型，同时尽可能降低过程中产生的成本。

硬件环境

        本次我们在开发验证阶段，采用的是推理卡，具体型号为310P3，在客户侧具体部署实施时，采用的是910B卡，本文介绍的技术方案在上述两种卡上均支持。

模型介绍

        本次涉及的其中一个模型是文本情感分类模型，大概结构是CNN层加上一些双向LSTM层，模型框架采用的是TensorFlow 1.x版本训练，之前的部署采用的是TFServing完成。

适配过程

方案一——原生TensorFlow + TF Adapter

        为了实现现有模型在NPU上运行，我们首先尝试了采用原生的TensorFlow + TFserving + TF Adapter 方案。按照官方提供的指导文档，安装1.x版本的TensorFlow和TFserving，选择从源码编译，加入NPU支持，具体可以参考链接。

        在完成环境准备后，启动TFServing进行模型推理部署，过程中发现报错信息，具体提示为CudnnLSTM算子不支持，无法完成自动转换。仔细查看模型定义后发现，该模型在训练定义和训练时用到了Keras中的一个与cuda相关的算子，即tf.contrib.cudnn_rnn.CudnnLSTM，该算子在Ascend提供的标准实现中不存在，导致转换和部署失败。

方案二——ATC

        注意到，Ascend提供了一个功能强大的工具ATC，可以将PyTorch或TensorFlow模型转换成offline model 即OM模型，再采用C++或Python接口进行推理。在完成环境准备后，采用以下命令进行模型转换：

atc --model=mymodel.pb \
    --framework=3 \
    --input_shape="input_1:1,30" \
    --out_nodes="output_node0:0" \
    --soc_version=Ascend310P3 \
    --display_model_info=1 \
    --log="debug" \
    --input_format="ND" \
    --output=save_model_om \
    --dynamic_batch_size="16,32" \

其中，mymodel.pb为模型参数文件。运行后发现，报错信息仍然存在，即ATC仍然不能将包换CudnnLSTM算子的模型转成OM格式

方案三——自定义Ascend C算子

        查询Ascend算子文档发现，目前提供的涉及LSTM的算子只有最简单的实现，没有与CudnnLSTM匹配的候选项，于是我们决定采用自定义Ascend C算子的方案，来实现模型的部署运行。

        首先，查阅TensorFlow文档，得到CudnnLSTM算子的定义如下：

tf.contrib.cudnn_rnn.CudnnLSTM(
    num_layers, num_units, input_mode=CUDNN_INPUT_LINEAR_MODE,
    direction=CUDNN_RNN_UNIDIRECTION, dropout=0.0, seed=None,
    dtype=tf.dtypes.float32, kernel_initializer=None, bias_initializer=None,
    name=None
)

其中，num_layers为LSTM层数，num_units为隐藏层维度。仔细查看该算子的推理代码后发现，其大致功能与简单的LSTM层并没有太大区别，因此实现难度不大。

        1 定义算子的输入输出，采用json文件指定：

[
    {
        "op": "CudnnLSTM",
         "input_desc": [
            {
                "name": "inputs",
                "param_type": "required",
                "format": ["ND"],
                "type": ["float32"]
            },
            {
                "name": "states_in",
                "param_type": "required",
                "format": ["ND"],
                "type": ["float32"]
            }
        ],
        "output_desc": [
            {
                "name": "hidden_states",
                "param_type": "required",
                "format": ["ND"],
                "type": ["float32"]
            },
            {
                "name": "states_out",
                "param_type": "required",
                "format": ["ND"],
                "type": ["float32"]
            },
        ]
    }
]

        2 按照自定义算子的流程，创建Ascend C算子工程：

${INSTALL_DIR}/python/site-packages/bin/msopgen gen -i cudnnlstm.json -c ai_core-310P3 -lan cpp -out CudnnLSTM

        3 Kernel侧算子实现

        在算子工程目标下，找到op_kernel/cudnnlstm.cpp文件，实现算子的核函数。该函数主要来实现LSTM模型的计算。根据LSTM的推倒公式，完成函数体。

        4 Host侧算子实现

         Host侧算子实现包括Tiling实现，shape推倒函数实现，算子的原型注册等。具体可以参阅官方提供的开发文档。该部分流程均比较固定，不涉及太多与LSTM算子相关的功能，因此不做赘述。

        5 算子工程编译

        在完成kenel和host开发后，可以将算子工程进行编译，生成.run安装包，从而实现部署运行。

        6 模型转换和部署

        在完成算子包的安装注册之后，再进行TFServing或者ATC转换，发现再也没有报错信息，一切正常。实测与cuda的结果对比，误差均在精度允许的范围内。

总结

        在遇到模型中存在某个Ascend中没有的算子时，通常可以采用自定义的方案。通过查阅框架中算子的原始代码和计算逻辑，将其翻译为Ascend C的开发语言，完成算子功能后进行编译、部署，即可增加对该算子的支持，从而实现在NPU卡运行任意模型。