An Overview of TVM and Model Optimization && TE

最新推荐文章于 2024-03-17 23:08:30 发布

pkapkaever

最新推荐文章于 2024-03-17 23:08:30 发布

阅读量218

点赞数

文章标签： python

本文链接：https://blog.csdn.net/pkapkaever/article/details/121111188

版权

TVM模型转化所采取的步骤

1.从框架导入模型，如果遇到问题，可以将其转换为ONNX
2.转换为Relay，这是TVM的high-level model language.会对模型进行一些图优化
3.lower to TE（张量表达式）。Relay运行的 FuseOps pass会将model划分为许多小的子图，将子图lower to TE表示。TE还提供了多个调度原语（such as tiling, vectorization, parallelization, unrolling, and fusion）来制定低级循环优化。（为了帮助将 Relay 表示转换为 TE 表示的过程，TVM 包括一个张量运算符清单 (TOPI)，它具有常见张量运算符（例如，conv2d、转置）的预定义模板。）
4.利用自动调整模块Auto TVM和Auto Schedule搜索最佳schedule。
AutoTVM：基于模板的自动调整模块。它运行搜索算法以在用户定义的模板中找到可调旋钮的最佳值。对于普通运营商，TOPI中已经提供了他们的模板。
AutoScheduler（又名 Ansor）：一个无模板的自动调整模块。它不需要预定义的计划模板。相反，它通过分析计算定义自动生成搜索空间。然后在生成的搜索空间中搜索最佳时间表。
5.选择模型编译的最佳配置。调优后，自动调优模块生成JSON格式的调优记录。此步骤为每个子图选择最佳时间表
6.lower to Tensor Intermediate Representation (TIR)，TVM 的低级中间表示。根据调整步骤选择最佳配置后，每个 TE 子图被降低到 TIR 并通过低级优化通道进行优化。接下来，优化的 TIR 被降低到硬件平台的目标编译器。这是生成可部署到生产中的优化模型的最终代码生成阶段。TVM 支持多种不同的编译器后端，包括：

LLVM，可以针对任意微处理器架构，包括标准 x86 和 ARM 处理器、AMDGPU 和 NVPTX 代码生成，以及 LLVM 支持的任何其他平台。

专用编译器，例如NVCC，NVIDIA 的编译器。

嵌入式和专用目标，通过 TVM 的自带代码生成 (BYOC) 框架实现。

使用TE手动优化矩阵乘法

先来一个没有进行手动优化的矩阵乘法，记录运行时间，方便对比

import tvm
import tvm.testing
from tvm import te
import numpy
import timeit
# The size of the matrix
# (M, K) x (K, N)
M = 1024
K = 1024
N = 1024
dtype = "float32"
target = tvm.target.Target(target="llvm", host="llvm")
dev = tvm.device(target.kind.name, 0)


# TVM Matrix Multiplication using TE
k = te.reduce_axis((0, K), "k")
A = te.placeholder((M, K), name="A")
B = te.placeholder((K, N), name="B")
C = te.compute((M, N), lambda x, y: te.sum(A[x, k] * B[k, y], axis=k), name="C")

# Default schedule
s = te.create_schedule(C.op)
a = tvm.nd.array(numpy.random.rand(M, K).astype(dtype), dev)
b = tvm.nd.array(numpy.random.rand(K, N).astype(dtype), dev)
c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)

#这是为了测试运算是否正确
# func = tvm.build(s, [A, B, C], target=target, name="mmult")
# func(a, b, c)
# answer = numpy.dot(a.asnumpy(), b.asnumpy())
# tvm.testing.assert_allclose(c.asnumpy(), answer, rtol=1e-5)


def evaluate_operation(s, vars, target, name, optimization, log):
    func = tvm.build(s, [A, B, C], target=target, name="mmult")
    assert func

    c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), dev)
    func(a, b, c)
    #tvm.testing.assert_allclose(c.asnumpy(), answer, rtol=1e-5)

    evaluator = func.time_evaluator(func.entry_name, dev, number=10)
    mean_time = evaluator(a, b, c).mean
    print("%s: %f" % (optimization, mean_time))
    log.append((optimization, mean_time))


log = []

evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="none", log=log)

print(tvm.lower(s, [A, B, C], simple_mode=True))

运行结果
在这里插入图片描述

1.Blocking

3232做平铺tile（内层循环，变成3232的小块），再使用因子factor=4做拆分（4条线程，32*32的小块，一共有256个）

####################################################1.blocking
bn = 32
# Blocking by loop tiling
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) = s[C].op.reduce_axis
#使用因子4做拆分（往后和矢量化操作有关，我感觉是可以多线程来的？）
ko, ki = s[C].split(k, factor=4)
# Hoist reduction domain outside the blocking loop
s[C].reorder(xo, yo, ko, ki, xi, yi)
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="blocking", log=log)
print(tvm.lower(s, [A, B, C], simple_mode=True))

运行结果
在这里插入图片描述
其中
我们来修改一些bn的值，可以看到这个tile，是讲C[1024][1024]根据axis[0]/16和axis[1]/32进行的划分

2.Vectorization
当内存访问模式一致时，编译器可以检测到这种模式并将连续内存传递给 SIMD 向量处理器。在 TVM 中，我们可以使用 vectorize接口来提示编译器这个模式，利用这个硬件特性

# Apply the vectorization optimization
s[C].vectorize(yi)
#在这里对yi矢量化
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="vectorization", log=log)

# The generalized IR after vectorization
print(tvm.lower(s, [A, B, C], simple_mode=True))

在这里插入图片描述

内层y.inner，没有了，可以看到变成了1，32

3.reorder

s = te.create_schedule(C.op)
xo,yo,xi,yi = s[C].tile(C.op.axis[0],C.op.axis[1],16,32)
(k,) = s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=8)
#可以看到这里交换了xi和ki的值
s[C].reorder(xo, yo, ko, xi ,ki, yi)
print("reorder:\n")
evaluate_operation(s, [A, B, C], target=target, name="reorder", optimization="reorder", log=log)
print(tvm.lower(s, [A, B, C], simple_mode=True))

在这里插入图片描述
4.数组打包

packedB = te.compute((N / bn, K, bn), lambda x, y, z: B[y, x * bn + z], name="packedB")
C = te.compute(
    (M, N),
    lambda x, y: te.sum(A[x, k] * packedB[y // bn, k, tvm.tir.indexmod(y, bn)], axis=k),
    name="C",
)

s = te.create_schedule(C.op)

xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) = s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)

s[C].reorder(xo, yo, ko, xi, ki, yi)
s[C].vectorize(yi)

x, y, z = s[packedB].op.axis
s[packedB].vectorize(z)
s[packedB].parallel(x)

evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="array packing", log=log)

# Here is the generated IR after array packing.
print(tvm.lower(s, [A, B, C], simple_mode=True))

运行结果
在这里插入图片描述

通过缓存块优化写入

s = te.create_schedule(C.op)
#分配写缓存
CC = s.cache_write(C,"global")  

xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
s[CC].compute_at(s[C],yo)

#新的inner轴
xc,yc = s[CC].op.axis
(k,) = s[CC].op.reduce_axis
ko, ki = s[CC].split(k, factor=4)

s[CC].reorder( ko, xc, ki, yc)
s[CC].unroll(ki)
s[CC].vectorize(yc)


# x, y, z = s[packedB].op.axis
# s[packedB].vectorize(z)
# s[packedB].parallel(x)

evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="block caching", log=log)

# Here is the generated IR after array packing.
print(tvm.lower(s, [A, B, C], simple_mode=True))