TVM Reduction降低算力
TVM Reduction降低算力
这是有关如何降低算力TVM的介绍材料。像sum / max / min这样的关联约简运算符是线性代数运算的典型构造块。
本文将演示如何降低TVM算力。
from __future__ import absolute_import, print_function
import tvm
import tvm.testing
from tvm import te
import numpy as np
描述行数
假设要计算行总数作为示例。用numpy语义可以写成B = numpy.sum(A, axis=1)
以下几行描述了行求和算子。创建归约公式，使用 te.reduce_axis来声明归约轴。te.reduce_axis降低算力的范围。 te.sum接受要降低算力的表达式以及降低算力轴，并计算声明范围内所有k的值之和。
等效的C代码如下：
for (int i = 0; i < n; ++i) {
B[i] = 0;
for (int k = 0; k < m; ++k) {
B[i] = B[i] + A[i][k];
}
}
n = te.var("n")
m = te.var("m")
A = te.placeholder((n, m), name="A")
k = te.reduce_axis((0, m), "k")
B = te.compute((n,), lambda i: te.sum(A[i, k], axis=k), name="B")
调度降低算力
有几种调度降低算力的方法。在执行任何操作之前，打印出默认调度的IR代码。
s = te.create_schedule(B.op)
print(tvm.lower(s, [A, B], simple_mode=True))
输出：
primfn(A_1: handle, B_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {B: Buffer(B_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
A: Buffer(A_2: Pointer(float32), float32, [n, m: int32], [stride_1: int32, stride_2: int32], type="auto")}
buffer_map = {A_1: A, B_1: B} {
for (i: int32, 0, n) {
B_2[(i*stride)] = 0f32
for (k: int32, 0, m) {
B_2[(i*stride)] = ((float32*)B_2[(i*stride)] + (float32*)A_2[((i*stride_1) + (k*stride_2))])
}
}
}
会发现IR代码与C代码非常相似。减速轴类似于法线轴，可以拆分。
在下面的代码中，将B的行轴和轴拆分为不同的因子。结果是嵌套归约。
ko, ki = s[B].split(B.op.reduce_axis[0], factor=16)
xo, xi = s[B].split(B.op.axis[0], factor=32)
print(tvm.lower(s, [A, B], simple_mode=True))
输出：
primfn(A_1: handle, B_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {B: Buffer(B_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
A: Buffer(A_2: Pointer(float32), float32, [n, m: int32], [stride_1: int32, stride_2: int32], type="auto")}
buffer_map = {A_1: A, B_1: B} {
for (i.outer: int32, 0, floordiv((n + 31), 32)) {
for (i.inner: int32, 0, 32) {
if @tir.likely((((i.outer*32) + i.inner) < n), dtype=bool) {
B_2[(((i.outer*32) + i.inner)*stride)] = 0f32
}
if @tir.likely((((i.outer*32) + i.inner) < n), dtype=bool) {
for (k.outer: int32, 0, floordiv((m + 15), 16)) {
for (k.inner: int32, 0, 16) {
if @tir.likely((((k.outer*16) + k.inner) < m), dtype=bool) {
B_2[(((i.outer*32) + i.inner)*stride)] = ((float32*)B_2[(((i.outer*32) + i.inner)*stride)] + (float32*)A_2[((((i.outer*32) + i.inner)*stride_1) + (((k.outer*16) + k.inner)*stride_2))])
}
}
}
}
}
}
}
要构建GPU内核，可以将B的行绑定到GPU线程。
s[B].bind(xo, te.thread_axis("blockIdx.x"))
s[B].bind(xi, te.thread_axis("threadIdx.x"))
print(tvm.lower(s, [A, B], simple_mode=True))
输出：
primfn(A_1: handle, B_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {B: Buffer(B_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
A: Buffer(A_2: Pointer(float32), float32, [n, m: int32], [stride_1: int32, stride_2: int32], type="auto")}
buffer_map = {A_1: A, B_1: B} {
attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = floordiv((n + 31), 32);
attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 32 {
if @tir.likely((((blockIdx.x*32) + threadIdx.x) < n), dtype=bool) {
B_2[(((blockIdx.x*32) + threadIdx.x)*stride)] = 0f32
}
for (k.outer: int32, 0, floordiv((m + 15), 16)) {
for (k.inner: int32, 0, 16) {
if @tir.likely((((blockIdx.x*32) + threadIdx.x) < n), dtype=bool) {
if @tir.likely((((k.outer*16) + k.inner) < m), dtype=bool) {
B_2[(((blockIdx.x*32) + threadIdx.x)*stride)] = ((float32*)B_2[(((blockIdx.x*32) + threadIdx.x)*stride)] + (float32*)A_2[((((blockIdx.x*32) + threadIdx.x)*stride_1) + (((k.outer*16) + k.inner)*stride_2))])
}
}
}
}
}
}
归约分解和并行化
建立归约的一个问题是，不能简单地在归约轴上并行化。需要对约简的算子进行划分，在对临时数组进行约简之前，将局部约简结果存储在临时数组中。
rfactor原语会重写计算。在下面的调度中，将B的结果写入临时结果B.rf。分解后的尺寸成为B.rf的第一尺寸。
s = te.create_schedule(B.op)
ko, ki = s[B].split(B.op.reduce_axis[0], factor=16)
BF = s.rfactor(B, ki)
print(tvm.lower(s, [A, B], simple_mode=True))
输出：
primfn(A_1: handle, B_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {B: Buffer(B_2: Pointer(float32), float32, [n: int32], [stride: int32], type="auto"),
A: Buffer(A_2: Pointer(float32), float32, [n, m: int32], [stride_1: int32, stride_2: int32], type="auto")}
buffer_map = {A_1: A, B_1: B} {
attr [B.rf: Pointer(float32)] "storage_scope" = "global";
allocate(B.rf, float32, [(n*16)]) {
for (k.inner: int32, 0, 16) {
for (i: int32, 0, n) {
B.rf[((k.inner*n) + i)] = 0f32
for (k.outer: int32, 0, floordiv((m + 15), 16)) {
if @tir.likely((((k.outer*16) + k.inner) < m), dtype=bool) {
B.rf[((k.inner*n) + i)] = ((float32*)B.rf[((k.inner*n) + i)] + (float32*)A_2[((i*stride_1) + (((k.outer*16) + k.inner)*stride_2))])
}
}
}
}
for (ax0: int32, 0, n) {
B_2[(ax0*stride)] = 0f32
for (k.inner.v: int32, 0, 16) {
B_2[(ax0*stride)] = ((float32*)B_2[(ax0*stride)] + (float32*)B.rf[((k.inner.v*n) + ax0)])
}
}
}
}
B的调度算子也将被重写为Bf缩减结果的第一轴上的和
print(s[B].op.body)
输出：
[reduce(combiner=comm_reducer(result=[(x + y)], lhs=[x], rhs=[y], identity_element=[0f]), source=[B.rf[k.inner.v, ax0]], init=[], axis=[iter_var(k.inner.v, range(min=0, ext=16))], where=(bool)1, value_index=0)]
降低算力跨线
现在，我们可以在分解后的轴上进行并行化处理。在此，B的复位轴标记为螺纹。TVM将算力减少轴标记为线程，如果它是唯一的算力降低，则可以在设备中进行交叉线程。
分解后的情况确实如此。也可以直接在还原轴上计算BF。最终生成的内核将按blockIdx.x划分行，按threadIdx.x划分threadIdx.y列，最后对threadIdx.x进行跨线程缩减
xo, xi = s[B].split(s[B].op.axis[0], factor=32)
s[B].bind(xo, te.thread_axis("blockIdx.x"))
s[B].bind(xi, te.thread_axis("threadIdx.y"))
tx = te.thread_axis("threadIdx.x")
s[B].bind(s[B].op.reduce_axis[0], tx)
s[BF].compute_at(s[B], s[B].op.reduce_axis[0])
s[B].set_store_predicate(tx.var.equal(0))
fcuda = tvm.build(s, [A, B], "cuda")
print(fcuda.imported_modules[0].get_source())
输出：
extern "C" __global__ void default_function_kernel0(float* __restrict__ A, float* __restrict__ B, int m, int n, int stride, int stride1, int stride2) {
float B_rf[1];
__shared__ float red_buf0[512];
B_rf[(0)] = 0.000000e+00f;
for (int k_outer = 0; k_outer < (m >> 4); ++k_outer) {
if (((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) < n) {
B_rf[(0)] = (B_rf[(0)] + A[(((((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) * stride) + (((k_outer * 16) + ((int)threadIdx.x)) * stride1)))]);
}
}
for (int k_outer1 = 0; k_outer1 < (((m & 15) + 15) >> 4); ++k_outer1) {
if (((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) < n) {
if (((((m >> 4) * 16) + (k_outer1 * 16)) + ((int)threadIdx.x)) < m) {
B_rf[(0)] = (B_rf[(0)] + A[(((((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) * stride) + (((((m >> 4) * 16) + (k_outer1 * 16)) + ((int)threadIdx.x)) * stride1)))]);
}
}
}
__syncthreads();
((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] = B_rf[(0)];
__syncthreads();
if (((int)threadIdx.x) < 8) {
((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] = (((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] + ((volatile float*)red_buf0)[((((((int)threadIdx.y) * 16) + ((int)threadIdx.x)) + 8))]);
((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] = (((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] + ((volatile float*)red_buf0)[((((((int)threadIdx.y) * 16) + ((int)threadIdx.x)) + 4))]);
((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] = (((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] + ((volatile float*)red_buf0)[((((((int)threadIdx.y) * 16) + ((int)threadIdx.x)) + 2))]);
((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] = (((volatile float*)red_buf0)[(((((int)threadIdx.y) * 16) + ((int)threadIdx.x)))] + ((volatile float*)red_buf0)[((((((int)threadIdx.y) * 16) + ((int)threadIdx.x)) + 1))]);
}
__syncthreads();
if (((int)threadIdx.x) == 0) {
B[((((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) * stride2))] = ((volatile float*)red_buf0)[((((int)threadIdx.y) * 16))];
}
}
将结果内核与numpy进行比较，验证结果内核的正确性。
nn = 128
ctx = tvm.gpu(0)
a = tvm.nd.array(np.random.uniform(size=(nn, nn)).astype(A.dtype), ctx)
b = tvm.nd.array(np.zeros(nn, dtype=B.dtype), ctx)
fcuda(a, b)
tvm.testing.assert_allclose(b.asnumpy(), np.sum(a.asnumpy(), axis=1), rtol=1e-4)
通过2D简化描述卷积
在TVM中，可以通过2D约简来描述卷积。这是2D卷积的示例，滤波器大小= [3，3]，步幅= [1，1]。
n = te.var("n")
Input = te.placeholder((n, n), name="Input")
Filter = te.placeholder((3, 3), name="Filter")
di = te.reduce_axis((0, 3), name="di")
dj = te.reduce_axis((0, 3), name="dj")
Output = te.compute(
(n - 2, n - 2),
lambda i, j: te.sum(Input[i + di, j + dj] * Filter[di, dj], axis=[di, dj]),
name="Output",
)
s = te.create_schedule(Output.op)
print(tvm.lower(s, [Input, Filter, Output], simple_mode=True))
出：
primfn(Input_1: handle, Filter_1: handle, Output_1: handle) -> ()
attr = {"global_symbol": "main", "tir.noalias": True}
buffers = {Output: Buffer(Output_2: Pointer(float32), float32, [(n: int32 - 2), (n - 2)], []),
Filter: Buffer(Filter_2: Pointer(float32), float32, [3, 3], []),
Input: Buffer(Input_2: Pointer(float32), float32, [n, n], [stride: int32, stride_1: int32], type="auto")}
buffer_map = {Input_1: Input, Filter_1: Filter, Output_1: Output} {
for (i: int32, 0, (n - 2)) {
for (j: int32, 0, (n - 2)) {
Output_2[((i*(n - 2)) + j)] = 0f32
for (di: int32, 0, 3) {
for (dj: int32, 0, 3) {
Output_2[((i*(n - 2)) + j)] = ((float32*)Output_2[((i*(n - 2)) + j)] + ((float32*)Input_2[(((i + di)*stride) + ((j + dj)*stride_1))]*(float32*)Filter_2[((di*3) + dj)]))
}
}
}
}
}
定义通用换向归约运算
除了内置的如降低算力操作te.sum， tvm.te.min和tvm.te.max，还可以通过定义交换降低算力操作te.comm_reducer。
n = te.var("n")
m = te.var("m")
product = te.comm_reducer(lambda x, y: x * y, lambda t: tvm.tir.const(1, dtype=t), name="product")
A = te.placeholder((n, m), name="A")
k = te.reduce_axis((0, m), name="k")
B = te.compute((n,), lambda i: product(A[i, k], axis=k), name="B")
注意
执行涉及多个值的归约argmax，可以通过元组输入来完成。有关更多详细信息，请参见使用协作输入来描述缩减。
总结
本文提供了降低算力调度的演练。
- 用reduce_axis描述归约。
- 如果需要并行性，请使用rfactor分解轴。
- 定义新的归约运算 te.comm_reducer
TVM Reduction降低算力的更多相关文章
- 端到端TVM编译器（上）
端到端TVM编译器(上) 摘要 将机器学习引入到各种各样的硬件设备中.AI框架依赖于特定于供应商的算子库,针对窄范围的服务器级gpu进行优化.将工作负载部署到新平台,例如手机.嵌入式设备和加速器(例如 ...
- CNN更新换代！性能提升算力减半，还即插即用
传统的卷积运算,要成为过去时了. Facebook和新加坡国立大学联手提出了新一代替代品:OctConv(Octave Convolution),效果惊艳,用起来还非常方便. OctConv就如同卷积 ...
- 创新全球算力生态价值，SPC算力生态强势来袭！
当前,区块链技术已经到了一个新的时代,即3.0时代.在区块链3.0时代,区块链技术迎来了数字经济革命,各行各业也在积极寻找与区块链能够融合的切入点.而随着区块链的愈加成熟,区块链技术也愈加被更多的人应 ...
- 如何使用TensorCores优化卷积
如何使用TensorCores优化卷积 本文将演示如何在TVM中使用TensorCores编写高性能的卷积计划.假设卷积的输入有大量数据.首先介绍如何在GPU上优化卷积. TensorCore简介 每 ...
- Inception V1、V2、V3和V4
Inception模块分为V1.V2.V3和V4. V1(GoogLeNet)的介绍 论文:Going deeper with convolutions 论文链接:https://arxiv.org/ ...
- 《区块链100问》第82集：应用类项目Golem
Golem是第一个基于以太坊区块链打造的计算资源交易平台.通过区块链,Golem能链接全球的算力资源,从而实现计算能力的全球共享.应用所有者和个体用户(算力“请求方”)可以点对点地从其他用户处租用算力 ...
- [转帖]中国AI芯“觉醒”的五年
中国AI芯“觉醒”的五年 https://www.cnbeta.com/articles/tech/857863.htm 原来 海思的营收已经超过了按摩店(AMD) 没想到.. 十多款芯片问世,多起并 ...
- 论文阅读笔记（二十三）【ECCV2018】：Robust Anchor Embedding for Unsupervised Video Person Re-Identification in the Wild
Introduction 当前主要的非监督方法都采用相同的训练数据集,这些数据集在不同摄像头中是对称的,即不存在单个行人的错误项,这些方法将在实际场景中效果下降.在本方法中,作者引入了非对称数据,如下 ...
- 自定义pass编写
自定义pass编写 TVM是一个框架,抽象了机器学习加速器的异质性.有时,用户可能需要自定义一些分析和IR转换,使TVM适应自己的专用硬件.本文可帮助用户在TVM中编写自定义pass. 先决条件 ...
随机推荐
- hdu4496并查集的删边操作
题意: 给你一个图,问你删除一些边后还有几个连通快.. 思路: 典型的并查集删边操作,并查集的删边就是先把不删除的边并查集一边(本题没有不删除的边),然后逆序吧所有要删除的边以 ...
- hdu3313 最大流找关键点,或者最短路找关键点.
题意: 给你一个有向图,然后给你起点和终点,问你从起点到终点有多少个关键点,如果当前的这个点删除了就无法从起点到终点,那么这个点就是一个关键点.. 思路: (1)有两种做法,我用的 ...
- Python脚本破解图形验证码(tesserocr和pytesseract)
在学习之前,我们先了解OCR.tesseract.tesserocr.pytesseract和opencv这几个跟图片处理有关的库. OCR(Optical Character Recognition ...
- 利用DNS进行命令控制和搭建隧道
目录 利用DNS进行命令控制(DNS-Shell) 利用DNS搭建隧道 利用DNS进行命令控制(DNS-Shell) DNS-Shell是一款通过DNS信道实现交互式Shell的强大工具,该工具的服务 ...
- 【easycode】使用说明 & 示例
语法 ## 是模板的注释,不会生成 # 全局的一些设置 $ 使用库函数或者调用全局定义的内容 常用方法 数据库表信息 ${tableInfo} comment 表备注 name 表名 fullColu ...
- Day009 二维数组
多维数组 多维数组是数组的嵌套(数组的元素是数组,数组的数组元素的元素是数组...),比如二维数组就是一个特殊的一维数组,其每一个元素都是一个一维数组. 二维数组 int a[][]=new int ...
- PowerDesigner16安装和使用
安装 安装参考链接:PowerDesigner安装教程 因为这个博主已经操作的很详细了,这边就不做过多的赘述. 使用 新建模型 选择物理模型 调出面板Palette 建表 最终的效果(一般不在数据库层 ...
- Wampserver-删除虚拟主机
对hosts操作 到目录C:\Windows\System32\drivers\etc中修改hosts 比如你想删除iwh2.com 选中这2行进行删除,保存退出 对httpd-vhosts操作 到目 ...
- Django（1）初识Django
前言 Django是一个开放源代码的Web应用框架,由Python写成,最初用于管理劳伦斯出版集团旗下的一些以新闻内容为主的网站,即CMS(内容管理系统)软件,于2005年7月在BSD许可证下发布,这 ...
- 软件测试中的测试用例Test Case原来是这么回事！
如果你去找一份功能测试的工作,在软件测试工程师面试过程中,有一些面试官会来一两个非常简单的问题 什么是Test Case?你是如何去写Test Case的? 我们先来看一下测试用例的介绍 什么是测试用 ...