初探TVM--使用Tensor Engine来编写算子

2023-11-22 07:38

本文主要是介绍初探TVM--使用Tensor Engine来编写算子,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

使用TE编写CPU算子

  • 什么是TE
  • 第一个例子:用TE写一个CPU的向量加法
    • 描述一个向量计算
    • 给计算创造一个默认的优化调度
    • 编译并且评估默认的调度
    • 使用并行化优化调度(paralleism)
    • 使用向量化的优化调度
    • 对比几种优化调度

什么是TE

就像在题目中写的那样,TE就是Tensor Engine的简称,其实就是用这些接口来定义一个计算算子是干什么的。可以转到tvm的一个大体介绍里面再看看。

通过这个教程,我们会学习TVM中,怎样通过TE定义一个张量计算算子。TE使用纯粹的功能性的语言描述一个张量计算,在这里是TE是没有边际效应(其实这个side effect一直不好理解,我觉得就是说换了硬件平台的话,也不应该对它产生影响)的。从tvm整体上看,relay是用来描述一组算子组成的计算图,TE是用来描述计算图中的每个算子节点,也就是说TE中可以认为是接受一个或多个输入张量,然后输出一个或多个输出张量。

在这个TE的入门的学习文章中,TVM用几个领域专用的张量计算来完成高效的算子构建。这里会通过两个算子来展示使用TE的技巧。第一个例子是个向量加法的例子,通过它来理解使用TE和TE的调度。然后我们会把这些调度概念一步一步拓展到一个compute bound的矩阵乘算子的例子上面去。同时,矩阵乘会被用到其他的TVM的操作当中,这个例子是一个非常好的入门的例子。

第一个例子:用TE写一个CPU的向量加法

我们先在Python的TE里面实现一个简单的向量加法吧,在中间加入一些基于CPU优化的优化调度试试。写Python的时候,还是老样子,首先import需要的依赖库。

import tvm 
import tvm.testing
import numpy as np
from tvm import te

还要再提一个说了三四次的东西了,就是如果把CPU的型号信息告知llvm,它会做出更多基于CPU特性的优化,比如被Linus疯狂吐槽的avx512。可以通过llc --version来查看CPU的版本信息,也可以通过cat /proc/cpuinfo来查看本台CPU的拓展新能包。例如,再查看过后可以使用llvm -mcpu=skylake-avx512来使能你的向量指令集。很遗憾,我用的AMD CPU,没有avx512. 不过CPU的型号是znver2,好像更高级一些啊。

tgt = tvm.target.Target(target="llvm", host="llvm")

描述一个向量计算

描述一个向量加法的计算。tvm支持向量语义,可以把每一层级的中间结果表达成为多维数组形式。用户需要描述生成向量的计算规则。我们可以先定义一个变量n来表示向量的形状,然后定义两个张量(tensor)占位符,AB,他们都有相同的形状(n,)。然后我们可以用一个compute操作来描述结果张量ccompute操作定义了一种计算,使得输出遵循一个特定的张量形状,并且在张量的每个元素上都做一个使用lambda表达式定义的计算。因为n此时是一个变量,所以A BC被定义了一个一致的形状。此时还没有实际的计算发生,我们仅仅声明了计算是什么。

n = te.var('n')
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda i: A[i]+B[i], name="C")

lambda表达式(lambda functions)
te.compute方法的第二个参数是一个执行计算的函数,在这个例子中间,我们使用了一个匿名函数的方式,也就是一个lambda表达式,本例中是期望对于AB的第i个元素做加法。

给计算创造一个默认的优化调度

虽然上面的几行代码定义了计算规则,我们仍然可以用不同的方式计算C以适应不同的设备。对于一个有多个维度的张量,你可以选择比如:1.首先计算的维度;2.哪些计算可以被拆分到不同的线程中。TVM要求使用者提供一套调度来描述计算是怎样被执行的。TE中的调度操作可以改变循环顺序,拆分计算到不同线程,以及对数据分块。一个重要的概念是,调度仅仅描述计算怎样被执行,也就是说,一个tensor expression上面应用不同的调度,应该产生同样的计算结果。

TVM允许创造一个基本的调度以row major的方式去计算C

for (int i = 0; i < n; i++)
{C[i] = A[i] + B[i];
}
s = te.create_schedule(C.op)

编译并且评估默认的调度

使用TE的表达式和一个优化调度,我们可以生成一个在目标架构和语言的可执行代码,现在就是LLVM和X86 CPU。我们给TVM提供:调度,TE表达式的列表,目标机器和主机,函数的名称。这个输出结果是一个无类型的函数,并且可以在python端直接被调用。

用下面一行代码,我们可以通过tvm.build生成一个函数。build函数拿到这个调度,函数签名(包含输入输出张量),以及目标语言。

fadd = tvm.build(s, [A, B, C], tgt, name="myadd")

我们可以运行这个函数,并且把输出结果和用numpy写的代码的结果作比较。编译后的TVM函数暴露出一个简明的C接口,这个结果可以通过任何语言调用。首先我们需要定义一个device,TVM可以在这个device上编译出相应的优化调度。当前我们的device是LLVM CPU。然后我们可以在device上初始化张量并且执行这个加法操作。我们通过对比TVM函数的输出结果和numpy的输出结果来验证正确性。

dev = tvm.device(tgt.kind.name, 0)n = 1024
a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)
c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)
fadd(a, b, c)
tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())

可以通过一个helper函数来profile这个TVM生成的代码,以此来和numpy函数比较一下运行速度。

import timeitnp_repeat = 100
np_running_time = timeit.timeit(setup="import numpy\n""n = 32768\n"'dtype = "float32"\n'"a = numpy.random.rand(n, 1).astype(dtype)\n""b = numpy.random.rand(n, 1).astype(dtype)\n",stmt="answer = a + b",number=np_repeat,
)
print("Numpy running time: %f" % (np_running_time / np_repeat))def evaluate_addition(func, target, optimization, log):dev = tvm.device(target.kind.name, 0)n = 32768a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)evaluator = func.time_evaluator(func.entry_name, dev, number=10)mean_time = evaluator(a, b, c).meanprint("%s: %f" % (optimization, mean_time))log.append((optimization, mean_time))log = [("numpy", np_running_time / np_repeat)]
evaluate_addition(fadd, tgt, "naive", log=log)

输出:

Numpy running time: 0.000010
naive: 0.000010

不得不说,还挺慢。。。再把架构指定一下吧:

tgt = tvm.target.Target(target="llvm -mcpu=znver2", host="llvm")

这次快了点儿:

Numpy running time: 0.000007
naive: 0.000005

使用并行化优化调度(paralleism)

我们随后开始看一看TE的基础内容,我们深入了解下调度是做什么的,以及怎么在目标平台上优化一个张量表达式。调度是一系列对于张量表达式的变形,当我们应用一个调度时,表达式的输入输出不变,仅仅在编译过程中的实现方法会改变。在张量加法算法中,加法是顺序被执行,因此比较容易使用处理器的多线程做并行化优化。代码这样子的:

s[C].parallel(C.op.axis[0])

tvm.lower()函数可以生成TE在对应优化调度下的IR(intermediate representation)。通过lower表达式的方式,我们可以查看不同调度对计算顺序的影响。我们可以用一个标志位simple_mode=true来输出一个可读的类C的伪代码:

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

会看到下面的输出:

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}buffers = {C: Buffer(C_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),A: Buffer(A_2: Pointer(float32), float32, [n], [stride_1: int32], type="auto"),B: Buffer(B_2: Pointer(float32), float32, [n], [stride_2: int32], type="auto")}buffer_map = {A_1: A, B_1: B, C_1: C} {for (i: int32, 0, n) "parallel" {C_2[(i*stride)] = ((float32*)A_2[(i*stride_1)] + (float32*)B_2[(i*stride_2)])}
}

现在就可以在不同的线程上并行运行这段表达式了,重新编译后可以再运行下这个并行化的表达式,看看效果:

fadd_parallel = tvm.build(s, [A, B, C], tgt, name="myadd_parallel")
fadd_parallel(a, b, c)tvm.testing.assert_allclose(c.numpy(), a.numpy(), b.numpy())evaluate_additiong(fadd_parallel, tgt, "parallel", log=log)

看看效果:

parallel: 0.000032

好像是比以前快了,但是没有到预期,因为128线程,只加速了两倍。猜测是因为n太小, 主要的耗时在访存上了,因此多线程加速起不到应有的效果。

使用向量化的优化调度

现在的各种CPU DSP都支持SIMD的向量化操作,我们可以用一个调度,把SIMD用上的。在TVM中,我们需要两个步骤来完成simd的调度:1.把循环拆分成内层和外层,内层是可以用单一simd指令完成的,外层是用于多线程完成的。由于要拆分循环,所以需要设定一个拆分的单位,以适配CPU的线程数和simd的宽度。

# Recreate the schedule, since we modified it with the parallel operation in
# the previous example
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")s = te.create_schedule(C.op)# This factor should be chosen to match the number of threads appropriate for
# your CPU. This will vary depending on architecture, but a good rule is
# setting this factor to equal the number of available CPU cores.
factor = 4outer, inner = s[C].split(C.op.axis[0], factor=factor)
s[C].parallel(outer)
s[C].vectorize(inner)fadd_vector = tvm.build(s, [A, B, C], tgt, name="myadd_parallel")evaluate_addition(fadd_vector, tgt, "vector", log=log)print(tvm.lower(s, [A, B, C], simple_mode=True))

看一下vectorize后的伪代码吧:

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}buffers = {C: Buffer(C_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),A: Buffer(A_2: Pointer(float32), float32, [n], [stride_1: int32], type="auto"),B: Buffer(B_2: Pointer(float32), float32, [n], [stride_2: int32], type="auto")}buffer_map = {A_1: A, B_1: B, C_1: C} {for (i.outer: int32, 0, floordiv((n + 3), 4)) "parallel" {for (i.inner.s: int32, 0, 4) {if @tir.likely((((i.outer*4) + i.inner.s) < n), dtype=bool) {C_2[(((i.outer*4) + i.inner.s)*stride)] = ((float32*)A_2[(((i.outer*4) + i.inner.s)*stride_1)] + (float32*)B_2[(((i.outer*4) + i.inner.s)*stride_2)])}}}
}

这时的耗时为:

vector: 0.000062

比不用simd还慢了。。。。

对比几种优化调度

上面尝试了很多调度后,我们可以对比一下:

baseline = log[0][1]
print("%s\t%s\t%s" % ("Operator".rjust(20), "Timing".rjust(20), "Performance".rjust(20)))
for result in log:print("%s\t%s\t%s"% (result[0].rjust(20), str(result[1]).rjust(20), str(result[1] / baseline).rjust(20)))
Operator                  Timing             Performancenumpy    6.184579997352557e-06                    1.0naive    4.635100000000001e-06     0.7494607559420615parallel             3.20865e-05       5.188145357281386vector    6.325680000000001e-05      10.22814807587232

我们在定义A/B/C的时候,给了他们一个相同的shape,就是n。tvm会利用这个相同的shape的特性做出些有效的优化。但是在调用优化库的时候,tvm会自动生成shape检测代码,如果给传入的向量不是这个shape,就会报错。也有些其他的方法来避免这些错误的出现,比如在定义运算的阶段,我们使用n=tvm.runtime.convert(1024)来代替n=te.var("n"),这样编译出的函数就会仅仅对1024内的长度做向量化。

使用tvm,我们能够定义、优化、编译一个向量加法的算子,并且在tvm的运行时组件上可以运行。这个算子可以被保存为一个库,在tvm的运行时组件中加载他。

这篇关于初探TVM--使用Tensor Engine来编写算子的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/408550

相关文章

Java Spring 中 @PostConstruct 注解使用原理及常见场景

《JavaSpring中@PostConstruct注解使用原理及常见场景》在JavaSpring中,@PostConstruct注解是一个非常实用的功能,它允许开发者在Spring容器完全初... 目录一、@PostConstruct 注解概述二、@PostConstruct 注解的基本使用2.1 基本代

C#使用StackExchange.Redis实现分布式锁的两种方式介绍

《C#使用StackExchange.Redis实现分布式锁的两种方式介绍》分布式锁在集群的架构中发挥着重要的作用,:本文主要介绍C#使用StackExchange.Redis实现分布式锁的... 目录自定义分布式锁获取锁释放锁自动续期StackExchange.Redis分布式锁获取锁释放锁自动续期分布式

springboot使用Scheduling实现动态增删启停定时任务教程

《springboot使用Scheduling实现动态增删启停定时任务教程》:本文主要介绍springboot使用Scheduling实现动态增删启停定时任务教程,具有很好的参考价值,希望对大家有... 目录1、配置定时任务需要的线程池2、创建ScheduledFuture的包装类3、注册定时任务,增加、删

使用Python实现矢量路径的压缩、解压与可视化

《使用Python实现矢量路径的压缩、解压与可视化》在图形设计和Web开发中,矢量路径数据的高效存储与传输至关重要,本文将通过一个Python示例,展示如何将复杂的矢量路径命令序列压缩为JSON格式,... 目录引言核心功能概述1. 路径命令解析2. 路径数据压缩3. 路径数据解压4. 可视化代码实现详解1

Pandas透视表(Pivot Table)的具体使用

《Pandas透视表(PivotTable)的具体使用》透视表用于在数据分析和处理过程中进行数据重塑和汇总,本文就来介绍一下Pandas透视表(PivotTable)的具体使用,感兴趣的可以了解一下... 目录前言什么是透视表?使用步骤1. 引入必要的库2. 读取数据3. 创建透视表4. 查看透视表总结前言

Python 交互式可视化的利器Bokeh的使用

《Python交互式可视化的利器Bokeh的使用》Bokeh是一个专注于Web端交互式数据可视化的Python库,本文主要介绍了Python交互式可视化的利器Bokeh的使用,具有一定的参考价值,感... 目录1. Bokeh 简介1.1 为什么选择 Bokeh1.2 安装与环境配置2. Bokeh 基础2

Android使用ImageView.ScaleType实现图片的缩放与裁剪功能

《Android使用ImageView.ScaleType实现图片的缩放与裁剪功能》ImageView是最常用的控件之一,它用于展示各种类型的图片,为了能够根据需求调整图片的显示效果,Android提... 目录什么是 ImageView.ScaleType?FIT_XYFIT_STARTFIT_CENTE

Java学习手册之Filter和Listener使用方法

《Java学习手册之Filter和Listener使用方法》:本文主要介绍Java学习手册之Filter和Listener使用方法的相关资料,Filter是一种拦截器,可以在请求到达Servl... 目录一、Filter(过滤器)1. Filter 的工作原理2. Filter 的配置与使用二、Listen

Pandas使用AdaBoost进行分类的实现

《Pandas使用AdaBoost进行分类的实现》Pandas和AdaBoost分类算法,可以高效地进行数据预处理和分类任务,本文主要介绍了Pandas使用AdaBoost进行分类的实现,具有一定的参... 目录什么是 AdaBoost?使用 AdaBoost 的步骤安装必要的库步骤一:数据准备步骤二:模型

使用Pandas进行均值填充的实现

《使用Pandas进行均值填充的实现》缺失数据(NaN值)是一个常见的问题,我们可以通过多种方法来处理缺失数据,其中一种常用的方法是均值填充,本文主要介绍了使用Pandas进行均值填充的实现,感兴趣的... 目录什么是均值填充?为什么选择均值填充?均值填充的步骤实际代码示例总结在数据分析和处理过程中,缺失数