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"
    }
  ],