TVM Compiler中文教程：TVM调度原语（Schedule Primitives）_tvm split-程序员宅基地

技术标签： split 调度Schedule tile TVM中文教程 TVM深度学习编译器

文章目录

TVM调度原语（Schedule Primitives）

TVM调度原语（Schedule Primitives）

TVM是用于高效内核代码构建的版本领域专用语言（Domain-Specialed-Language，DSL）。

这篇教程，我们将展示通过TVM提供的各种原语怎么去调度计算。

from __future__ import absolute_import, print_function

import tvm
import numpy as np

通常存在几种计算相同结果的方法，但是，不同的方法将导致不同的局部性和性能，所以TVM要求用户提供怎么去调用Schedule描述计算是如何执行的。

Schedule是计算的变换的集合，它通过变换程序中的计算循环Loop，实现不同性能。

#定义一些变量
n = tvm.var('n')
m = tvm.var('m')

调度（Schedule）能通过一系列计算ops来定义，默认情况下，调度计算张量以航为顺序。

#定义矩阵元素element-wise乘法
A = tvm.placeholder((m,n), name='A')
B = tvm.placeholder((m,n), name='B')
C = tvm.compute((m,n),lambda i,j: A[i,j] * B[i,j], name ='C')
#创建调度
s = tvm.create_schedule([C.op])
#lower会将计算从定义转换为真正的可调用函数。 使用参数`simple_mode = True`，它将返回一个可读的C伪代码，我们在这里使用它来打印计划结果。
print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C {
    
  for (i, 0, m) {
    
    for (j, 0, n) {
    
      C[((i*n) + j)] = (A[((i*n) + j)]*B[((i*n) + j)])
    }
  }
}

一个调度过程由多个阶段组成，一个阶段表示操作的一个调度。我们提供各种方法来调度每个阶段。

分裂split

split通过factor分裂指定轴为两个轴。

A = tvm.placeholder((m,), name='A')
B = tvm.compute((m,), lambda i: A[i]*2, name='B')

s = tvm.create_schedule(B.op)
#分裂0轴为两个轴，先计算内循环再计算外循环，xo为外循环，xi为内循环
xo, xi = s[B].split(B.op.axis[0], factor=32)
print(tvm.lower(s, [A, B], simple_mode=True))

输出：

produce B {
    
  for (i.outer, 0, ((m + 31)/32)) {
    
    for (i.inner, 0, 32) {
    
      if (likely(((i.outer*32) < (m - i.inner)))) {
    
        B[((i.outer*32) + i.inner)] = (A[((i.outer*32) + i.inner)]*2.000000f)
      }
    }
  }
}

使用nparts与factor作用相反，nparts指定外循环次数，factor指定内循环次数。

A = tvm.placeholder((m,), name='A')
B = tvm.compute((m,), lambda i: A[i], name='B')

s = tvm.create_schedule(B.op)
bx, tx = s[B].split(B.op.axis[0], nparts=32)
print(tvm.lower(s, [A, B], simple_mode=True))

输出：

produce B {
    
  for (i.outer, 0, 32) {
    
    for (i.inner, 0, ((m + 31)/32)) {
    
      if (likely((i.inner < (m - (i.outer*((m + 31)/32)))))) {
    
        if (likely(((0 - (i.outer*((m + 31)/32))) <= i.inner))) {
    
          B[(i.inner + (i.outer*((m + 31)/32)))] = A[(i.inner + (i.outer*((m + 31)/32)))]
        }
      }
    }
  }
}

平铺tile

tile通过平铺两个轴执行计算图块

A = tvm.placeholder((m, n), name='A')
B = tvm.compute((m, n), lambda i, j: A[i, j], name='B')

s = tvm.create_schedule(B.op)
xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)
print(tvm.lower(s, [A, B], simple_mode=True))

输出：

produce B {
    
  for (i.outer, 0, ((m + 9)/10)) {
    
    for (j.outer, 0, ((n + 4)/5)) {
    
        //先执行10x5的图块，滑动下一个图块
      for (i.inner, 0, 10) {
    
        for (j.inner, 0, 5) {
    
          if (likely(((i.outer*10) < (m - i.inner)))) {
    
            if (likely(((j.outer*5) < (n - j.inner)))) {
    
              B[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)] = A[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)]
            }
          }
        }
      }
    }
  }
}

融合fuse

fuse能融合一个计算的两个轴

A = tvm.placeholder((m,n),name='A')
B = tvm.compute((m,n), lambda i,j: A[i,j], name='B')

s = tvm.create_schedule(B.op)
#首先平铺成4轴（i.outer,j.outer,i.inner,j.inner）
xo,yo,xi,yi = s[B].tile(B.op.axis[0],B.op.axis[1], x_factor=10, y_factor=5)
#然后融合(i.inner,j.inner)进一个轴：（i.inner.j.inner.fused）
fused = s[B].fuse(xi,yj)
print(tvm.lower(s, [A, B], simple_mode=True))

输出：

produce B {
  for (i.outer, 0, ((m + 9)/10)) {
    for (j.outer, 0, ((n + 4)/5)) {
      for (i.inner.j.inner.fused, 0, 50) {
        if (likely(((i.outer*10) < (m - (i.inner.j.inner.fused/5))))) {
          if (likely(((j.outer*5) < (n - (i.inner.j.inner.fused % 5))))) {
            B[(((j.outer*5) + (i.inner.j.inner.fused % 5)) + (((i.outer*10) + (i.inner.j.inner.fused/5))*n))] = A[(((j.outer*5) + (i.inner.j.inner.fused % 5)) + (((i.outer*10) + (i.inner.j.inner.fused/5))*n))]
          }
        }
      }
    }
  }
}

重排序reorder

reorder能按照指定顺序重新排列轴（类似于permute）。

A = tvm.placeholder((m, n), name='A')
B = tvm.compute((m, n), lambda i, j: A[i, j], name='B')
s = tvm.create_schedule(B.op)
#首先平铺成4轴（i.outer,j.outer,i.inner,j.inner）
xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)
s[B].reorder(xi,yo,xo,yi)
print(tvm.lower(s, [A, B], simple_mode=True))

输出：

produce B {
    
  for (i.inner, 0, 10) {
    
    for (j.outer, 0, ((n + 4)/5)) {
    
      for (i.outer, 0, ((m + 9)/10)) {
    
        for (j.inner, 0, 5) {
    
          if (likely(((i.outer*10) < (m - i.inner)))) {
    
            if (likely(((j.outer*5) < (n - j.inner)))) {
    
              B[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)] = A[(((j.outer*5) + (((i.outer*10) + i.inner)*n)) + j.inner)]
            }
          }
        }
      }
    }
  }
}

绑定bind

bind可以使用线程轴绑定指定的轴，通常在GPU编程中使用。

A = tvm.placeholder((n,), name='A')
B = tvm.compute(A.shape, lambda i: A[i] * 2, name='B')

s = tvm.create_schedule(B.op)
bx, tx = s[B].split(B.op.axis[0], factor=64)
s[B].bind(bx, tvm.thread_axis("blockIdx.x"))
s[B].bind(tx, tvm.thread_axis("threadIdx.x"))
print(tvm.lower(s, [A, B], simple_mode=True))

输出：

produce B {
    
  // attr [iter_var(blockIdx.x, , blockIdx.x)] thread_extent = ((n + 63)/64)
  // attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 64
  if (likely(((blockIdx.x*64) < (n - threadIdx.x)))) {
    
    B[((blockIdx.x*64) + threadIdx.x)] = (A[((blockIdx.x*64) + threadIdx.x)]*2.000000f)
  }
}

从哪里开始计算compute_at

对于包含多个算子的调度，TVM默认从root开始遍历计算张量。

A = tvm.placeholder((m,), name='A')
B = tvm.compute((m,), lambda i: A[i]+1, name='B')
C = tvm.compute((m,), lambda i: B[i]*2, name='C')

s = tvm.create_schedule(C.op)
print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce B {
    
  for (i, 0, m) {
    
    B[i] = (A[i] + 1.000000f)
  }
}
produce C {
    
  for (i, 0, m) {
    
    C[i] = (B[i]*2.000000f)
  }
}

compute_at可以将B的计算移动到C的第一个计算轴。

A = tvm.placeholder((m,), name='A')
B = tvm.compute((m,), lambda i: A[i]+1, name='B')
C = tvm.compute((m,), lambda i: B[i]*2, name='C')

s = tvm.create_schedule(C.op)
# 移动B循环到C循环的第一个轴
s[B].compute_at(S[C], C.op.axis[0])
print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C {
    
  for (i, 0, m) {
    
    produce B {
    
      B[i] = (A[i] + 1.000000f)
    }
    C[i] = (B[i]*2.000000f)
  }
}

计算内联compute_inline

compute_inline可以将一个计算阶段标记为内联，然后将计算体扩展并插入需要张量的地址处。（和C中的内联函数一个意思）

A = tvm.placeholder((m,), name='A')
B = tvm.compute((m,), lambda i: A[i]+1, name='B')
C = tvm.compute((m,), lambda i: B[i]*2, name='C')

s = tvm.create_schedule(C.op)
s[B].compute_inline()
print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C {
    
  for (i, 0, m) {
    
    //类似内联函数，直接合成一个循环
    C[i] = ((A[i]*2.000000f) + 2.000000f)
  }
}

compute_root

compute_root可以将一个计算阶段的计算移动到root。（compute_at的逆过程）

A = tvm.placeholder((m,), name='A')
B = tvm.compute((m,), lambda i: A[i]+1, name='B')
C = tvm.compute((m,), lambda i: B[i]*2, name='C')

s = tvm.create_schedule(C.op)
# B移动到C的0轴
s[B].compute_at(s[C], C.op.axis[0])
# B重新移动回root
s[B].compute_root()
print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce B {
    
  for (i, 0, m) {
    
    B[i] = (A[i] + 1.000000f)
  }
}
produce C {
    
  for (i, 0, m) {
    
    C[i] = (B[i]*2.000000f)
  }
}

总结

本教程介绍了tvm中的调度原语，它允许用户轻松灵活地调度计算。

为了获得良好性能的内核实现，一般工作流程通常是：

通过一系列操作描述您的计算。
尝试使用原语安排计算。
编译并运行以查看性能差异。
根据运行结果调整你的调度。

本文链接：https://blog.csdn.net/hw5226349/article/details/91370350

原作者删帖不实内容删帖广告或垃圾文章投诉

智能推荐

使用JDBC连接数据库出现 The server time zone value ‘�й��׼ʱ��‘ is unrecognized or represents more than one解决方案_jdbc.properties timezone-程序员宅基地

文章浏览阅读553次。在 jdbc.properties 文件中的 url 后面加上 ?serverTimezone=UTC加入之前的jdbc.properties文件：user=rootpassword=12345678url=jdbc:mysql://localhost:3306/testdriverClass=com.mysql.cj.jdbc.Driver加入之后：user=rootpassword=12345678url=jdbc:mysql://localhost:3306/test?serv_jdbc.properties timezone

计算机图形学孔令德基础知识,计算机图形学基础教程孔令德答案-程序员宅基地

文章浏览阅读1.4k次。计算机图形学基础教程孔令德答案【篇一：大学计算机图形学课程设】息科学与工程学院课程设计任务书题目：小组成员：巴春华、焦国栋成员学号：专业班级：计算机科学与技术、2009级本2班课程：计算机图形学指导教师：燕孝飞职称：讲师完成时间： 2011年12 月----2011年 12 月枣庄学院信息科学与工程学院制2011年12 月20日课程设计任务书及成绩评定12【篇二：计算机动画】第一篇《计算机图形学》..._计算机图形学基础教程孔令德答案

python xlwings追加数据_大数据分析Python库xlwings提升Excel工作效率教程-程序员宅基地

文章浏览阅读1k次。原标题：大数据分析Python库xlwings提升Excel工作效率教程Excel在当今的企业中非常非常普遍。在AAA教育，我们通常建议出于很多原因使用代码，并且我们的许多数据科学课程旨在教授数据分析和数据科学的有效编码。但是，无论您偏爱使用大数据分析Python的程度如何，最终，有时都需要使用Excel来展示您的发现或共享数据。但这并不意味着仍然无法享受大数据分析Python的某些效率!实际上，..._xlwings通过索引添加数据

java8u211_jre864位u211-程序员宅基地

文章浏览阅读911次。iefans为用户提供的jre8 64位是针对64位windows平台而开发的java运行环境软件，全称为java se runtime environment 8，包括Java虚拟机、Java核心类库和支持文件，不包含开发工具--编译器、调试器和其它工具。jre需要辅助软件--JavaPlug-in--以便在浏览器中运行applet。本次小编带来的是jre8 64位官方版下载，版本小号u211版..._jre8是什么

kasp技术原理_KASP基因分型-程序员宅基地

文章浏览阅读5k次。KASP基因分型介绍KASP(Kompetitive Allele-Specific PCR)，即竞争性等位基因特异性PCR，原理上与TaqMan检测法类似，都是基于终端荧光信号的读取判断，每孔反应都是采用双色荧光检测一个SNP位点的两种基因型，不同的SNP对应着不同的荧光信号。KASP技术与TaqMan法类似，它与TaqMan技术不同的是，它不需要每个SNP位点都合成特异的荧光引物，它基于独特的..._kasp是什么

华为p50预装鸿蒙系统,华为p50会不会预装鸿蒙系统_华为p50会预装鸿蒙系统吗-程序员宅基地

文章浏览阅读154次。华为现在比较火的还真就是新开发的鸿蒙系统了，那么在即将上市的华为p50手机上会不会预装鸿蒙系统呢？接下来我们就来一起了解一下华为官方发布的最新消息吧。1.华为p50最新消息相信大家都知道，随着华为鸿蒙OS系统转正日期临近，似乎全网的花粉们都在关注华为鸿蒙OS系统优化、生态建设等等，直接忽略了不断延期发布的华为P50手机，如今华为P50系列手机终于传来了最新的好消息，在经过一系列方案修改以后，终于被..._华为手机p50直接预装鸿蒙系统