初探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使用Curator进行ZooKeeper操作的详细教程

《Java使用Curator进行ZooKeeper操作的详细教程》ApacheCurator是一个基于ZooKeeper的Java客户端库,它极大地简化了使用ZooKeeper的开发工作,在分布式系统... 目录1、简述2、核心功能2.1 CuratorFramework2.2 Recipes3、示例实践3

springboot security使用jwt认证方式

《springbootsecurity使用jwt认证方式》:本文主要介绍springbootsecurity使用jwt认证方式,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地... 目录前言代码示例依赖定义mapper定义用户信息的实体beansecurity相关的类提供登录接口测试提供一

go中空接口的具体使用

《go中空接口的具体使用》空接口是一种特殊的接口类型,它不包含任何方法,本文主要介绍了go中空接口的具体使用,具有一定的参考价值,感兴趣的可以了解一下... 目录接口-空接口1. 什么是空接口?2. 如何使用空接口?第一,第二,第三,3. 空接口几个要注意的坑坑1:坑2:坑3:接口-空接口1. 什么是空接

springboot security快速使用示例详解

《springbootsecurity快速使用示例详解》:本文主要介绍springbootsecurity快速使用示例,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝... 目录创www.chinasem.cn建spring boot项目生成脚手架配置依赖接口示例代码项目结构启用s

Python如何使用__slots__实现节省内存和性能优化

《Python如何使用__slots__实现节省内存和性能优化》你有想过,一个小小的__slots__能让你的Python类内存消耗直接减半吗,没错,今天咱们要聊的就是这个让人眼前一亮的技巧,感兴趣的... 目录背景:内存吃得满满的类__slots__:你的内存管理小助手举个大概的例子:看看效果如何?1.

java中使用POI生成Excel并导出过程

《java中使用POI生成Excel并导出过程》:本文主要介绍java中使用POI生成Excel并导出过程,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教... 目录需求说明及实现方式需求完成通用代码版本1版本2结果展示type参数为atype参数为b总结注:本文章中代码均为

Spring Boot3虚拟线程的使用步骤详解

《SpringBoot3虚拟线程的使用步骤详解》虚拟线程是Java19中引入的一个新特性,旨在通过简化线程管理来提升应用程序的并发性能,:本文主要介绍SpringBoot3虚拟线程的使用步骤,... 目录问题根源分析解决方案验证验证实验实验1:未启用keep-alive实验2:启用keep-alive扩展建

使用Java实现通用树形结构构建工具类

《使用Java实现通用树形结构构建工具类》这篇文章主要为大家详细介绍了如何使用Java实现通用树形结构构建工具类,文中的示例代码讲解详细,感兴趣的小伙伴可以跟随小编一起学习一下... 目录完整代码一、设计思想与核心功能二、核心实现原理1. 数据结构准备阶段2. 循环依赖检测算法3. 树形结构构建4. 搜索子

GORM中Model和Table的区别及使用

《GORM中Model和Table的区别及使用》Model和Table是两种与数据库表交互的核心方法,但它们的用途和行为存在著差异,本文主要介绍了GORM中Model和Table的区别及使用,具有一... 目录1. Model 的作用与特点1.1 核心用途1.2 行为特点1.3 示例China编程代码2. Tab

SpringBoot使用OkHttp完成高效网络请求详解

《SpringBoot使用OkHttp完成高效网络请求详解》OkHttp是一个高效的HTTP客户端,支持同步和异步请求,且具备自动处理cookie、缓存和连接池等高级功能,下面我们来看看SpringB... 目录一、OkHttp 简介二、在 Spring Boot 中集成 OkHttp三、封装 OkHttp