初探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

相关文章

中文分词jieba库的使用与实景应用(一)

知识星球:https://articles.zsxq.com/id_fxvgc803qmr2.html 目录 一.定义: 精确模式(默认模式): 全模式: 搜索引擎模式: paddle 模式(基于深度学习的分词模式): 二 自定义词典 三.文本解析   调整词出现的频率 四. 关键词提取 A. 基于TF-IDF算法的关键词提取 B. 基于TextRank算法的关键词提取

使用SecondaryNameNode恢复NameNode的数据

1)需求: NameNode进程挂了并且存储的数据也丢失了,如何恢复NameNode 此种方式恢复的数据可能存在小部分数据的丢失。 2)故障模拟 (1)kill -9 NameNode进程 [lytfly@hadoop102 current]$ kill -9 19886 (2)删除NameNode存储的数据(/opt/module/hadoop-3.1.4/data/tmp/dfs/na

Hadoop数据压缩使用介绍

一、压缩原则 (1)运算密集型的Job,少用压缩 (2)IO密集型的Job,多用压缩 二、压缩算法比较 三、压缩位置选择 四、压缩参数配置 1)为了支持多种压缩/解压缩算法,Hadoop引入了编码/解码器 2)要在Hadoop中启用压缩,可以配置如下参数

Makefile简明使用教程

文章目录 规则makefile文件的基本语法:加在命令前的特殊符号:.PHONY伪目标: Makefilev1 直观写法v2 加上中间过程v3 伪目标v4 变量 make 选项-f-n-C Make 是一种流行的构建工具,常用于将源代码转换成可执行文件或者其他形式的输出文件(如库文件、文档等)。Make 可以自动化地执行编译、链接等一系列操作。 规则 makefile文件

使用opencv优化图片(画面变清晰)

文章目录 需求影响照片清晰度的因素 实现降噪测试代码 锐化空间锐化Unsharp Masking频率域锐化对比测试 对比度增强常用算法对比测试 需求 对图像进行优化,使其看起来更清晰,同时保持尺寸不变,通常涉及到图像处理技术如锐化、降噪、对比度增强等 影响照片清晰度的因素 影响照片清晰度的因素有很多,主要可以从以下几个方面来分析 1. 拍摄设备 相机传感器:相机传

pdfmake生成pdf的使用

实际项目中有时会有根据填写的表单数据或者其他格式的数据,将数据自动填充到pdf文件中根据固定模板生成pdf文件的需求 文章目录 利用pdfmake生成pdf文件1.下载安装pdfmake第三方包2.封装生成pdf文件的共用配置3.生成pdf文件的文件模板内容4.调用方法生成pdf 利用pdfmake生成pdf文件 1.下载安装pdfmake第三方包 npm i pdfma

零基础学习Redis(10) -- zset类型命令使用

zset是有序集合,内部除了存储元素外,还会存储一个score,存储在zset中的元素会按照score的大小升序排列,不同元素的score可以重复,score相同的元素会按照元素的字典序排列。 1. zset常用命令 1.1 zadd  zadd key [NX | XX] [GT | LT]   [CH] [INCR] score member [score member ...]

git使用的说明总结

Git使用说明 下载安装(下载地址) macOS: Git - Downloading macOS Windows: Git - Downloading Windows Linux/Unix: Git (git-scm.com) 创建新仓库 本地创建新仓库:创建新文件夹,进入文件夹目录,执行指令 git init ,用以创建新的git 克隆仓库 执行指令用以创建一个本地仓库的

【北交大信息所AI-Max2】使用方法

BJTU信息所集群AI_MAX2使用方法 使用的前提是预约到相应的算力卡,拥有登录权限的账号密码,一般为导师组共用一个。 有浏览器、ssh工具就可以。 1.新建集群Terminal 浏览器登陆10.126.62.75 (如果是1集群把75改成66) 交互式开发 执行器选Terminal 密码随便设一个(需记住) 工作空间:私有数据、全部文件 加速器选GeForce_RTX_2080_Ti

【Linux 从基础到进阶】Ansible自动化运维工具使用

Ansible自动化运维工具使用 Ansible 是一款开源的自动化运维工具,采用无代理架构(agentless),基于 SSH 连接进行管理,具有简单易用、灵活强大、可扩展性高等特点。它广泛用于服务器管理、应用部署、配置管理等任务。本文将介绍 Ansible 的安装、基本使用方法及一些实际运维场景中的应用,旨在帮助运维人员快速上手并熟练运用 Ansible。 1. Ansible的核心概念