运行和验证

TVM编译安装完成后进行TVM的运行和验证。

为了验证基于开源TVM合入softmax算子优化补丁以及集成毕昇编译器所带来的性能提升，本章节提供一个softmax算子测试脚本，对比在开源TVM 0.9.0环境和按照本文档安装的TVM环境下编译得到的softmax算子基于鲲鹏920新型号处理器上的运行性能。因此，为了完成性能对比，用户除按照本文档准备TVM环境之外，还需准备一套开源TVM 0.9.0环境，具体请参见TVM官方文档。

在“/home”目录下创建测试目录“/test”，并在此目录下编写softmax测试脚本test_softmax.py。

创建“test_softmax.py”文件。

cd /home
mkdir test
cd test
vi test_softmax.py

按“i”进入编辑模式，编写“test_softmax.py”文件，添加以下内容。

import pytest
import tvm
from tvm import relay
import numpy as np
from tvm.contrib import graph_executor
from tvm.contrib.debugger import debug_executor
import os
import argparse
os.environ["TVM_BIND_THREADS"] = "0"
def test_softmax(opt_level="3", repeat_time=50):
    # Define a larger and more complex shape and data for the softmax input
    shape = (100, 120, 50)  # Increase dimensions for greater complexity
    data = np.random.rand(*shape).astype(np.float32)  # Use random values for variability
    # Create a Relay expression for the softmax function
    data_var = relay.var("data", shape=shape, dtype="float32")
    softmax_expr = relay.nn.softmax(data_var, axis=-1)
    # Create a Relay function and module
    func = relay.Function([data_var], softmax_expr)
    mod = tvm.IRModule.from_expr(func)
    target = tvm.target.Target("llvm -mtriple=aarch64-linux-gnu -mattr=+neon")
    # Create and build the module
    with tvm.transform.PassContext(opt_level=int(3)):
        lib = relay.build(mod, target=target)
        try:
            print(lib.function_metadata["tvmgen_default_fused_nn_softmax"])
        except:
            print(lib.function_metadata["tvmgen_default_fused_nn_fast_softmax"])
    model_path = "/home/test/test_softmax_%s.so" % (opt_level)
    lib.export_library(model_path)
    lib = tvm.runtime.load_module(model_path)
    # Create TVM runtime
    dev = tvm.cpu()
    data = tvm.nd.array(data, device=dev)
    m = graph_executor.create(lib["get_graph_json"](), lib, dev)
    # Set input data
    m.set_input("data", data)
    # Run the module and get the output
    m.run()
    output = m.get_output(0).numpy()
    
    # Create debug executor (ensure `lib` and `dev` are correctly passed)
    m_debug = debug_executor.create(
        lib["get_graph_json"](), 
        lib, 
        dev, 
        dump_root="./tvmdbg_onnx"
    )
    
    m.set_input("data", data)
    results = [0] * repeat_time
    ops_sums = [0] * repeat_time
    for i in range(repeat_time):
        report = m_debug.profile(**{"data": data})
        r_t = report.table().split()
        for ind, tmp in enumerate(r_t):
            if tmp.endswith("softmax"):
                time_cons, time_port = r_t[ind + 1], r_t[ind + 2]
                time_cons = time_cons.replace(",", "")
                ops_sums[i] = float(time_cons)
                
            if tmp == "Total":
                # 保存总时间（推理总耗时）
                total_time = r_t[ind + 1]
                total_time = total_time.replace(",", "")
                results[i] = float(total_time)
    print("Average time consuming in softmaxOp is %f ms" % (sum(ops_sums) / repeat_time / 1000))
    print("Average time consuming in profiling is %f ms" % (sum(results) / repeat_time / 1000))
    # Validate the output
    data = data.numpy()
    expected_output = np.exp(data) / np.sum(np.exp(data), axis=-1, keepdims=True)
    np.testing.assert_allclose(output, expected_output, atol=1e-5)
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--num_threads", type=str, default="4")
    parser.add_argument("--opt_level", type=str, default="1")
    parser.add_argument("--epoch", type=int, default=100)
    parser.add_argument("--repeat_time", type=int, default=50)
    args = parser.parse_args()
    # 设置测试脚本的线程数
    num_threads = args.num_threads
    os.environ["TVM_NUM_THREADS"] = num_threads
    # 设置优化等级
    opt_level = args.opt_level
    # 设置推理次数
    epoch = args.epoch
    # 设置profile重复次数
    repeat_time = args.repeat_time
    print("num of threads is :%d" % (tvm.runtime.num_threads()))
    print("opt level is :%s" % opt_level)
    print("epoch is %d" % epoch)
    print("repeat time is %d" % repeat_time)
    
    # 验证softmax算子正确性
    test_softmax(opt_level=opt_level, repeat_time=repeat_time)
    pytest.main([__file__])

按“Esc”键，输入:wq!，按“Enter”保存并退出编辑。

激活TVM的Python环境，并设置TVM环境变量（激活开源TVM 0.9.0，需将“/path/to/TVM”替换为开源TVM 0.9.0的安装路径）。
1 2 3
conda activate tvm export TVM_HOME=/path/to/TVM/apache-tvm-src-v0.9.0 export PYTHONPATH=$TVM_HOME/python:${PYTHONPATH}
运行test_softmax.py测试脚本。
1
taskset -c 0-15 python -u test_softmax.py --num_threads 4
test_softmax.py脚本会打印TVM为softmax算子生成的底层中间表示，以及算子耗时。

如图1和图2所示测试脚本执行结果，执行结果分析如下。
- 对比softmax算子的底层中间表示。本文档指导安装TVM环境下，算子计算增加了向量化和并行化，并在最后的正则化运算循环中，将除法提取到循环外。
- 对比softmax算子耗时。本文档指导安装TVM环境下，编译得到的softmax算子耗时低于开源TVM 0.9.0环境。
图1 开源TVM 0.9.0环境中test_softmax.py运行结果

 图2 本文档指导安装的TVM环境中test_softmax.py运行结果