使用Tensorize评估硬件内部特性
使用Tensorize评估硬件内部特性
这是有关如何在TVM中执行张量的入门文档。
通过使用调度原语tensorize，人们可以用相应的内部函数代替计算单元，从而轻松利用handcrafted micro-kernels，扩展TVM以支持新的硬件体系结构。
本文的目的是展示张量的功能和用法，而不是提供有效的解决方案。
from __future__ import absolute_import, print_function
import tvm
from tvm import te
import numpy as np
定义矩阵乘法
以矩阵乘法为例。Matmul首先将两个矩阵之间的对应元素相乘，然后在某个轴上累积。以下几行描述了TVM中A * B^T的计算。
N, M, L = 1024, 512, 64
A = te.placeholder((N, L), name="A")
B = te.placeholder((M, L), name="B")
k = te.reduce_axis((0, L), name="k")
C = te.compute((N, M), lambda i, j: te.sum(A[i, k] * B[j, k], axis=k), name="C")
s = te.create_schedule(C.op)
print(tvm.lower(s, [A, B, C], simple_mode=True))
输出：
primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
B: Buffer(B_2: Pointer(float32), float32, [512, 64], []),
A: Buffer(A_2: Pointer(float32), float32, [1024, 64], [])}
buffer_map = {A_1: A, B_1: B, C_1: C} {
for (i: int32, 0, 1024) {
for (j: int32, 0, 512) {
C_2[((i*512) + j)] = 0f32
for (k: int32, 0, 64) {
C_2[((i*512) + j)] = ((float32*)C_2[((i*512) + j)] + ((float32*)A_2[((i*64) + k)]*(float32*)B_2[((j*64) + k)]))
}
}
}
}
调度Matmul
假设有一个支持矩阵矢量乘法（GEMV）作为硬件原语的加速器，可以采用任意大小的reduce轴，但另一个轴必须不大于16。因此，分解了matmul循环，生成最里面的一个（16x64）GEMV循环。
factor = 16
x, y = C.op.axis
(z,) = C.op.reduce_axis
yo, yi = s[C].split(y, factor=factor)
s[C].reorder(x, yo, yi, z)
print(tvm.lower(s, [A, B, C], simple_mode=True))
输出：
primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
B: Buffer(B_2: Pointer(float32), float32, [512, 64], []),
A: Buffer(A_2: Pointer(float32), float32, [1024, 64], [])}
buffer_map = {A_1: A, B_1: B, C_1: C} {
for (i: int32, 0, 1024) {
for (j.outer: int32, 0, 32) {
for (j.inner: int32, 0, 16) {
C_2[(((i*512) + (j.outer*16)) + j.inner)] = 0f32
for (k: int32, 0, 64) {
C_2[(((i*512) + (j.outer*16)) + j.inner)] = ((float32*)C_2[(((i*512) + (j.outer*16)) + j.inner)] + ((float32*)A_2[((i*64) + k)]*(float32*)B_2[(((j.outer*1024) + (j.inner*64)) + k)]))
}
}
}
}
}
如上面打印的IR所示，内部循环j.inner与k一起形成GEMV的计算-在最内部的两个循环内，索引i是固定的，对矩阵的访问A仅变化k，生成A“向量”的访问模式”。 可以用j.inner张量来评估假设的硬件的GEMV指令。
定义固有的GEMV张量化
调度张量前，先定义GEMV的固有函数。它包括两部分，第一部分是GEMV的计算定义。TVM使用它来匹配原始Matmul调度中的计算模式。第二个是指定如何在设备上执行GEMV，这在intrin_func下面完成。
def intrin_gemv(m, l):
a = te.placeholder((l,), name="a")
b = te.placeholder((m, l), name="b")
k = te.reduce_axis((0, l), name="k")
c = te.compute((m,), lambda i: te.sum(a[k] * b[i, k], axis=k), name="c")
Ab = tvm.tir.decl_buffer(a.shape, a.dtype, name="A", offset_factor=1, strides=[1])
Bb = tvm.tir.decl_buffer(b.shape, b.dtype, name="B", offset_factor=1, strides=[te.var("s1"), 1])
Cb = tvm.tir.decl_buffer(c.shape, c.dtype, name="C", offset_factor=1, strides=[1])
def intrin_func(ins, outs):
ib = tvm.tir.ir_builder.create()
aa, bb = ins
cc = outs[0]
ib.emit(
tvm.tir.call_extern(
"int32",
"gemv_update",
cc.access_ptr("w"),
aa.access_ptr("r"),
bb.access_ptr("r"),
m,
l,
bb.strides[0],
)
)
return ib.get()
return te.decl_tensor_intrin(c.op, intrin_func, binds={a: Ab, b: Bb, c: Cb})
在此te.decl_tensor_intrin声明如何执行计算c.op。实现只接受输入和输出，将它们转换为指针并发出外部函数调用。注意，张量需要用户指定offset_factor，原始数据结构的起始地址和传递给张量的偏移量之间对齐的问题，TVM能评估，通过矢量化加载进行优化。为了简化，将系数设置为1。
为输入和输出声明了缓冲区，尽管这不是必需的，将从缓冲区提供的额外信息中受益。例如，bb.strides[0]作为参数传递 给外部函数gemv_update。将看到bb.strides[0] == l如何与更复杂的调度区分开。
注意，将te.var("s1")用作第一个步幅B。如果可以推理出步幅（在这种情况下，TVM确定张量B是紧凑的，步幅是[L, 1]），可以使用此类placeholder让TVM自动为绑定推理的值。
gemv = intrin_gemv(factor, L)
s[C].tensorize(yi, gemv)
print(tvm.lower(s, [A, B, C], simple_mode=True))
输出：
primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
B: Buffer(B_2: Pointer(float32), float32, [512, 64], []),
A: Buffer(A_2: Pointer(float32), float32, [1024, 64], [])}
buffer_map = {A_1: A, B_1: B, C_1: C} {
for (i: int32, 0, 1024) {
for (j.outer: int32, 0, 32) {
@tir.call_extern("gemv_update", @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), C_2, ((i*512) + (j.outer*16)), 16, 2, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), A_2, (i*64), 64, 1, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), B_2, (j.outer*1024), 1024, 1, dtype=handle), 16, 64, 64, dtype=int32)
}
}
}
通过张大yi，最里面的两个循环被之前定义的内在函数代替。为了构建和运行该模块，定义外部函数gemv_update，它是GEMV的naive实现，仅用于演示。
def gemv_impl():
cc_code = """
extern "C" int gemv_update(float *cc, float *aa, float *bb, int m, int l, int stride) {
for (int i = 0; i < m; ++i) {
for (int j = 0; j < l; ++j) {
cc[i] += aa[j] * bb[i * stride + j];
}
}
return 0;
}
"""
from tvm.contrib import utils, clang
temp = utils.tempdir()
ll_path = temp.relpath("temp.ll")
# Create LLVM ir from c source code
ll_code = clang.create_llvm(cc_code, output=ll_path)
return ll_code
利用pragma属性import_llvm导入llvm asm内联。导入在执行张量的GEMV之前进行。
s[C].pragma(x, "import_llvm", gemv_impl())
print(tvm.lower(s, [A, B, C], simple_mode=True))
出：
primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
B: Buffer(B_2: Pointer(float32), float32, [512, 64], []),
A: Buffer(A_2: Pointer(float32), float32, [1024, 64], [])}
buffer_map = {A_1: A, B_1: B, C_1: C} {
attr [IterVar(i: int32, (nullptr), "DataPar", "")] "pragma_import_llvm" = "; ModuleID = '/tmp/tmpinr5hwkd/input0.cc'
source_filename = "/tmp/tmpinr5hwkd/input0.cc"
target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
target triple = "x86_64-pc-linux-gnu"
; Function Attrs: noinline nounwind optnone uwtable
define dso_local i32 @gemv_update(float*, float*, float*, i32, i32, i32) #0 {
%7 = alloca float*, align 8
%8 = alloca float*, align 8
%9 = alloca float*, align 8
%10 = alloca i32, align 4
%11 = alloca i32, align 4
%12 = alloca i32, align 4
%13 = alloca i32, align 4
%14 = alloca i32, align 4
store float* %0, float** %7, align 8
store float* %1, float** %8, align 8
store float* %2, float** %9, align 8
store i32 %3, i32* %10, align 4
store i32 %4, i32* %11, align 4
store i32 %5, i32* %12, align 4
store i32 0, i32* %13, align 4
br label %15
15: ; preds = %50, %6
%16 = load i32, i32* %13, align 4
%17 = load i32, i32* %10, align 4
%18 = icmp slt i32 %16, %17
br i1 %18, label %19, label %53
19: ; preds = %15
store i32 0, i32* %14, align 4
br label %20
20: ; preds = %46, %19
%21 = load i32, i32* %14, align 4
%22 = load i32, i32* %11, align 4
%23 = icmp slt i32 %21, %22
br i1 %23, label %24, label %49
24: ; preds = %20
%25 = load float*, float** %8, align 8
%26 = load i32, i32* %14, align 4
%27 = sext i32 %26 to i64
%28 = getelementptr inbounds float, float* %25, i64 %27
%29 = load float, float* %28, align 4
%30 = load float*, float** %9, align 8
%31 = load i32, i32* %13, align 4
%32 = load i32, i32* %12, align 4
%33 = mul nsw i32 %31, %32
%34 = load i32, i32* %14, align 4
%35 = add nsw i32 %33, %34
%36 = sext i32 %35 to i64
%37 = getelementptr inbounds float, float* %30, i64 %36
%38 = load float, float* %37, align 4
%39 = fmul float %29, %38
%40 = load float*, float** %7, align 8
%41 = load i32, i32* %13, align 4
%42 = sext i32 %41 to i64
%43 = getelementptr inbounds float, float* %40, i64 %42
%44 = load float, float* %43, align 4
%45 = fadd float %44, %39
store float %45, float* %43, align 4
br label %46
46: ; preds = %24
%47 = load i32, i32* %14, align 4
%48 = add nsw i32 %47, 1
store i32 %48, i32* %14, align 4
br label %20
49: ; preds = %20
br label %50
50: ; preds = %49
%51 = load i32, i32* %13, align 4
%52 = add nsw i32 %51, 1
store i32 %52, i32* %13, align 4
br label %15
53: ; preds = %15
ret i32 0
}
attributes #0 = { noinline nounwind optnone uwtable "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-frame-pointer-elim"="true" "no-frame-pointer-elim-non-leaf" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="false" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" }
!llvm.module.flags = !{!0}
!llvm.ident = !{!1}
!0 = !{i32 1, !"wchar_size", i32 4}
!1 = !{!"clang version 9.0.1-+20191211110317+c1a0a213378-1~exp1~20191211221711.104 "}
";
for (i, 0, 1024) {
for (j.outer: int32, 0, 32) {
@tir.call_extern("gemv_update", @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), C_2, ((i*512) + (j.outer*16)), 16, 2, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), A_2, (i*64), 64, 1, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), B_2, (j.outer*1024), 1024, 1, dtype=handle), 16, 64, 64, dtype=int32)
}
}
}
最后，将张量版本与numpy.dot产生的张量版本进行比较，确保实现正确。
func = tvm.build(s, [A, B, C], target="llvm", name="gemv")
from tvm.topi.utils import get_const_tuple
dtype = A.dtype
ctx = tvm.context("cpu", 0)
a = np.random.uniform(size=get_const_tuple(A.shape)).astype(dtype)
b = np.random.uniform(size=get_const_tuple(B.shape)).astype(dtype)
c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=dtype), ctx)
func(tvm.nd.array(a, ctx), tvm.nd.array(b, ctx), c)
tvm.testing.assert_allclose(c.asnumpy(), np.dot(a, b.T), rtol=1e-3)
进行Tensorize更新Reduce-update
已经了解了张量化的基本概念，现在让向更复杂的情况迈进一步。
假设加速器只能将向量乘以一个矩阵，向量的大小必须不大于16。考虑到硬件限制，需要按如下方式拆分reduce轴。
zo, zi = s[C].split(z, factor=factor)
s[C].reorder(x, yo, zo, yi, zi)
由于张量内在函数现在仅覆盖了reduce轴的一部分，而不是使用一个“ body”函数，因此TVM需要一个reduce_reset在reduce for循环之前调用的reduce_update函数，以及一个定义“ update”的函数。计算策略。
def gemv_impl():
cc_code = """
extern "C" int gemv_update(float *cc, float *aa, float *bb, int m, int l, int stride) {
for (int i = 0; i < m; ++i) {
for (int j = 0; j < l; ++j) {
cc[i] += aa[j] * bb[i * stride + j];
}
}
return 0;
}
extern "C" int gemv_reset(float *cc, int m) {
for (int i = 0; i < m; ++i) {
cc[i] = 0.0;
}
return 0;
}
"""
from tvm.contrib import utils, clang
temp = utils.tempdir()
ll_path = temp.relpath("temp.ll")
# Create LLVM ir from c source code
ll_code = clang.create_llvm(cc_code, output=ll_path)
return ll_code
def intrin_gemv(m, l):
a = te.placeholder((l,), name="a")
b = te.placeholder((m, l), name="b")
k = te.reduce_axis((0, l), name="k")
c = te.compute((m,), lambda i: te.sum(a[k] * b[i, k], axis=k), name="c")
Ab = tvm.tir.decl_buffer(a.shape, a.dtype, name="A", offset_factor=1, strides=[1])
Bb = tvm.tir.decl_buffer(b.shape, b.dtype, name="B", offset_factor=1, strides=[te.var("s1"), 1])
Cb = tvm.tir.decl_buffer(c.shape, c.dtype, name="C", offset_factor=1, strides=[1])
def intrin_func(ins, outs):
aa, bb = ins
cc = outs[0]
def _body():
ib = tvm.tir.ir_builder.create()
ib.emit(
tvm.tir.call_extern(
"int32",
"gemv_update",
cc.access_ptr("w"),
aa.access_ptr("r"),
bb.access_ptr("r"),
m,
l,
bb.strides[0],
)
)
return ib.get()
def _reduce_reset():
ib = tvm.tir.ir_builder.create()
ib.emit(tvm.tir.call_extern("int32", "gemv_reset", cc.access_ptr("w"), m))
return ib.get()
def _reduce_update():
return _body()
return _body(), _reduce_reset(), _reduce_update()
return te.decl_tensor_intrin(c.op, intrin_func, binds={a: Ab, b: Bb, c: Cb})
注意，intrin_func返回一个三元组： 如果tensorization包括所有的reduce轴，功能将被调用，否则一起将被使用。在示例中，共享相同的实现，而在其它情况下，硬件对于这两个功能可能具有不同的指令。此外，由于平铺，可以看到现在是不同的。(body, reduce_reset, reduce_update)body()reduce_reset()reduce_update()body()reduce_update()bb.strides[0]l
张量squared GEMV，生成并检查结果
gemv = intrin_gemv(factor, factor)
s[C].tensorize(yi, gemv)
s[C].pragma(yo, "import_llvm", gemv_impl())
func = tvm.build(s, [A, B, C], target="llvm", name="gemv")
a = np.random.uniform(size=get_const_tuple(A.shape)).astype(dtype)
b = np.random.uniform(size=get_const_tuple(B.shape)).astype(dtype)
c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=dtype), ctx)
func(tvm.nd.array(a, ctx), tvm.nd.array(b, ctx), c)
tvm.testing.assert_allclose(c.asnumpy(), np.dot(a, b.T), rtol=1e-3)
概要
本文演示了TVM中张量内在函数的用法。Tensorize为用户提供了一种通过微内核获得完全优化的调度方式。例如，英特尔CPU上使用张量化直接调用AVX指令进行INT8量化。使TVM可以编译为ASIC-有关详细信息，请参阅VTA：深度学习加速器堆栈。演示了如何使用内联程序集导入，这可以帮助用户轻松地将asm输入调度中。
使用Tensorize评估硬件内部特性的更多相关文章
- js的内部特性--属性
使用方法:通过调用Object.defineProperty(对象,"对象属性",{}进行的操作}) 当对一个对象的属性的属性类型中vlue设置为一个值时,则这个对象的这个属性的值 ...
- 【翻译】.NET 5中的性能改进
[翻译].NET 5中的性能改进 在.NET Core之前的版本中,其实已经在博客中介绍了在该版本中发现的重大性能改进. 从.NET Core 2.0到.NET Core 2.1到.NET Core ...
- 【翻译】.NET 5 Preview7发布
今天,发布了.NET 5.0 Preview7.这是倒数第二个预览版本(在转移到RC之前).此时,大多数功能应该已经非常接近完成了.Single file和ARM64 intrinsics是两个花费了 ...
- Linux转发性能评估与优化-转发瓶颈分析与解决方式(补遗)
补遗 关于网络接收的软中断负载均衡,已经有了成熟的方案,可是该方案并不特别适合数据包转发,它对server的小包处理非常好.这就是RPS.我针对RPS做了一个patch.提升了其转发效率. 下面是我转 ...
- Linux转发性能评估与优化(转发瓶颈分析与解决方式)
线速问题 非常多人对这个线速概念存在误解. 觉得所谓线速能力就是路由器/交换机就像一根网线一样. 而这,是不可能的.应该考虑到的一个概念就是延迟. 数据包进入路由器或者交换机,存在一个核心延迟操作,这 ...
- Flash硬件原理
1.2.1. 什么是Flash Flash全名叫做Flash Memory,从名字就能看出,是种数据存储设备,存储设备有很多类,Flash属于非易失性存储设备(Non-volatile Memory ...
- 深入GPU硬件架构及运行机制
目录 一.导言 1.1 为何要了解GPU? 1.2 内容要点 1.3 带着问题阅读 二.GPU概述 2.1 GPU是什么? 2.2 GPU历史 2.2.1 NV GPU发展史 2.2.2 NV GPU ...
- Microsoft Tech Summit 2018 课程简述：利用 Windows 新特性开发出更好的手绘视频应用
概述 Microsoft Tech Summit 2018 微软技术暨生态大会将于10月24日至27日在上海世博中心举行,这也会是国内举办的最后一届 Tech Summit,2019 年开始会以 Mi ...
- CSS硬件加速的好与坏
本文翻译自Ariya Hidayat的Hardware Accelerated CSS: The Nice vs The Naughty.感谢Kyle He帮助校对. 每个人都痴迷于60桢每秒的顺滑动 ...
随机推荐
- 【转】如何用MTR诊断网络问题
MTR 是一个强大的网络诊断工具,管理员能够用它诊断和隔离网络错误,并向上游提供商提供有关网络状态的有用报告.MTR 通过更大的采样来跟踪路由,就像 traceroute + ping 命令的组合.本 ...
- Android平台dalvik模式下java Hook框架ddi的分析（1）
本文博客地址:http://blog.csdn.net/qq1084283172/article/details/75710411 一.前 言 在前面的博客中已经学习了作者crmulliner编写的, ...
- Win64 驱动内核编程-31.枚举与删除映像回调
枚举与删除映像回调 映像回调可以拦截 RING3 和 RING0 的映像加载.某些游戏保护会用此来拦截黑名单中的驱动加载,比如 XUETR.WIN64AST 的驱动.同理,在反游戏保护的过程中,也可以 ...
- IDEA 这样设置，好看到爆炸！！！
Hello,大家好,我是楼下小黑哥. 今天这篇文章是次条视频的文案,这里推荐大家直接看视频学习. IDEA 这样设置,好看到爆炸!!!#01 今天这期我们来分享几个美化 IDEA 设置技巧,让你的 I ...
- 【Redis】redis异步消息队列+Spring自定义注解+AOP方式实现系统日志持久化
说明: SSM项目中的每一个请求都需要进行日志记录操作.一般操作做的思路是:使用springAOP思想,对指定的方法进行拦截.拼装日志信息实体,然后持久化到数据库中.可是仔细想一下会发现:每次的客户端 ...
- Morgan Stanley Books List：经典金融书籍推荐
一．经济学 1. 中华帝国的专制制度,佛朗索瓦.魁奈 2. 资本论(共3卷),马恩全集 3. 国家竞争优势,麦克尔.波特 4. Essentials of corporate analysis, by ...
- OOP第一章总结
经过了三周的OO,尽管过程不太轻松,但是有所得还是值得欣慰的事! (1)程序结构 第一次作业: UML类图如下,第一次作业在结构上并没有太多面向对象的思想,只是简单的分类,一个运行类,两个对象类,预处 ...
- .Net·快速查找哪一个类库引用了哪一个Package
阅文时长 | 0.18分钟 字数统计 | 348.8字符 主要内容 | 1.引言&背景 2.查找法示例 3.声明与参考资料 『.Net·快速查找哪一个类库引用了哪一个Package』 编写人 ...
- Swagger简明教程
一.什么是swagger Swagger 是一个规范和完整的框架,用于生成.描述.调用和可视化RESTful风格的 Web 服务.总体目标是使客户端和文件系统作为服务器以同样的速度来更新.文件的方法, ...
- [bug] Junit initializationError
原因 导包错误 解决 先删除 import org.junit.Test; 再导入正确的包 参考 https://blog.csdn.net/javae100/article/details/7978 ...