AICompiler是集成在PAI-Blade中的AI编译优化组件，包含Static Shape和Dynamic Shape编译框架。通常您无需提供额外配置，AICompiler即可在通用透明的情况下帮助您提高推理性能。本文介绍如何使用AICompiler对TensorFlow和PyTorch模型进行编译优化。

背景信息

随着AI模型结构的快速演化，底层计算硬件的层出不穷，用户使用习惯的推陈出新，使得单纯基于手工优化解决AI模型的性能和效率问题越来越困难。为了解决这些问题，AI编译优化技术已经成为一个被广泛关注的技术方向。

传统编译器是以高层语言作为输入，避免用户直接写机器码。深度学习编译器的作用与其相仿，输入是比较灵活的、具备较高抽象度的计算图，输出包括CPU或GPU等硬件平台上的底层机器码及执行引擎。AI编译器的目标是针对AI计算任务，以通用编译器的方式完成性能优化。让用户可以专注于上层模型开发，降低用户手工优化性能的人力开发成本，进一步压榨硬件性能空间。

在过去两年多时间里，PAI团队在AI编译优化技术方向投入了比较专注的资源精力，AICompiler已经作为优化组件之一集成到PAI-Blade之中，从而帮助用户以通用透明的方式完成推理模型的优化和部署。

目前AICompiler主要包含Static Shape和Dynamic Shape两套编译框架。Dynamic Shape可以支持所有任务类型，包括计算图Shape变化范围较大的推理任务。Static Shape主要适配计算图为静态Shape的任务，或Shape变化范围较小的任务，以获得理论上更加极致的性能。Static Shape和Dynamic Shape的关系如下表所示。
编译框架 TensorFlow PyTorch 适配范围
Static Shape 支持 暂不支持 适配静态Shape或Shape变化范围较小的任务，理论上可以获得极致的性能收益。
Dynamic Shape 支持 支持 适配所有的任务类型。
默认情况下，PAI-Blade会自动分析用户的模型是否适合Dynamic Shape的编译优化，用户无需提供额外的输入信息。对于静态Shape类任务或Shape变化范围较小的任务，用户可以提供额外的输入选项，此时PAI-Blade会尝试通过Static Shape模式编译以获得更优的性能。下文详细介绍几个应用例子供您参考。

通过Dynamic Shape模式编译TensorFlow模型

以一个开源ASR模型为例，演示如何通过Dynamic Shape模式编译TensorFlow模型：
  1. 下载模型和测试数据。
    # 下载示例模型、测试数据。
wget https://pai-blade.cn-hangzhou.oss.aliyun-inc.com/test_public_model/bbs/tf_aicompiler_demo/frozen.pb
wget https://pai-blade.cn-hangzhou.oss.aliyun-inc.com/test_public_model/bbs/tf_aicompiler_demo/test_bc4.npy
  2. 加载模型、指定数据，并调用blade.optimize，无需任何额外配置。
    import numpy as np
import tensorflow as tf
import blade

# 加载模型和测试数据。
graph_def = tf.GraphDef()
with open('./frozen.pb', 'rb') as f:
    graph_def.ParseFromString(f.read())
test_data = np.load('test_bc4.npy', allow_pickle=True, encoding='bytes',).item()

# 开始优化。
optimized_model, opt_spec, report = blade.optimize(
    graph_def,  # The original model, here is a TF GraphDef.
    'o1',  # Optimization level o1 or o2.
    device_type='gpu',  # Target device to run the optimized model.
    config=blade.Config(),
    inputs=['encoder_memory_placeholder', 'encoder_memory_length_placeholder'],
    outputs=['score', 'seq_id'],
    test_data=[test_data]
    #verbose=True
)

# 保存优化结果。
tf.train.write_graph(optimized_model, "./", "optimized.pb", as_text=False)
print("Report: {}".format(report))
  3. 优化完成后，查看优化报告（blade.optimize返回的report）中AICompiler生效之后的性能收益。本例中以T4卡为例，AICompiler可以帮您在通用透明的情况下获得2.23倍的性能收益，report中的字段含义请参见优化报告。
        {
      "name": "TfAicompilerGpu",
      "status": "effective",
      "speedup": "2.23",
      "pre_run": "120.54 ms",
      "post_run": "53.99 ms"
    }

通过Static Shape模式编译TensorFlow模型

如果模型部署后的实际Shape变化范围较小或完全是Static Shape，您可以通过配置config选项，指定PAI-Blade以Static Shape模式进行编译。相比Dynamic Shape模式编译，您只需修改blade.optimize函数的config入参即可实现Static Shape编译，示例代码如下。
optimized_model, opt_spec, report = blade.optimize(
    graph_def,  # The original model, here is a TF GraphDef.
    'o1',  # Optimization level o1 or o2.
    device_type='gpu',  # Target device to run the optimized model.
    # Provide an additional config here in order to try Static Shape Compilation:
    config=blade.Config(enable_static_shape_compilation_opt = True), 
    inputs=['encoder_memory_placeholder', 'encoder_memory_length_placeholder'],
    outputs=['score', 'seq_id'],
    test_data=[test_data]
    #verbose=True
)
关于config参数的高级配置，详情请参见表 1。
优化完成后，可以查看优化报告（blade.optimize返回的report）中的性能收益。本例中的性能收益进一步提升至2.35倍，report中的字段含义请参见优化报告。
    {
      "name": "TfAicompilerGpu",
      "status": "effective",
      "speedup": "2.35",
      "pre_run": "114.91 ms",
      "post_run": "48.86 ms"
    }

通过Dynamic Shape模式编译PyTorch模型

以一个开源ASR模型为例，演示如何通过Dynamic Shape模式编译PyTorch模型：
  1. 下载模型。
    # PyTorch 1.6.0
# Python3.6
wget https://pai-blade.cn-hangzhou.oss.aliyun-inc.com/test_public_model/bbs/pt_aicompiler_demo/orig_decoder_v2.pt
  2. 加载模型、指定数据，并调用blade.optimize，无需任何额外配置。
    import os
import time
import torch
# To use blade, just import it.
import blade

# 加载模型。
pt_file = 'orig_decoder_v2.pt'
batch = 8
model = torch.jit.load(pt_file)

# 准备测试数据。
def get_test_data(batch_size=1):
    decoder_input_t = torch.LongTensor([1] * batch_size).cuda()
    decoder_hidden_t = torch.rand(batch_size, 1, 256).cuda()
    decoder_hidden_t = decoder_hidden_t * 1.0
    decoder_hidden_t = torch.tanh(decoder_hidden_t)
    output_highfeature_t = torch.rand(batch_size, 448, 4, 50).cuda()
    attention_sum_t = torch.rand(batch_size, 1, 4, 50).cuda()
    decoder_attention_t = torch.rand(batch_size, 1, 4, 50).cuda()
    et_mask = torch.rand(batch_size, 4, 50).cuda()

    return (decoder_input_t, decoder_hidden_t, output_highfeature_t, attention_sum_t, decoder_attention_t, et_mask)

dummy = get_test_data(batch)

# 开始优化。
optimized_model, opt_spec, report = blade.optimize(
    model,  # The original model, here is a torch scrip model.
    'o1',  # Optimization level o1 or o2.
    device_type='gpu',  # Target device to run the optimized model.
    test_data=[dummy],  # For PyTorch, input data is list of tupoles.
    config=blade.Config()
)

print("spec: {}".format(opt_spec))
print("report: {}".format(report))

# 导出优化结果。
torch.jit.save(optimized_model, 'optimized_decoder.pt')
  3. 优化完成后，查看优化报告（blade.optimize返回的report）中AICompiler生效之后的性能收益。本例中以T4卡为例，AICompiler可以帮您在通用透明的情况下获得2.45倍的性能收益，report中的字段含义请参见优化报告。
      "optimizations": [
    {
      "name": "PyTorchMlir",
      "status": "effective",
      "speedup": "2.45",
      "pre_run": "1.99 ms",
      "post_run": "0.81 ms"
    }
  ],