TVM Compiler中文教程：TVM调度原语（Schedule Primitives）

最新推荐文章于 2024-06-13 23:13:46 发布
最新推荐文章于 2024-06-13 23:13:46 发布
阅读量8k 收藏 25
点赞数 6
分类专栏： TVM深度学习编译器 文章标签： TVM中文教程 调度Schedule split tile
本文是TVM中文教程的一部分，重点介绍TVM调度原语，包括split、tile、fuse、reorder、bind、compute_at、compute_inline和compute_root等，这些原语用于调整计算顺序和局部性，以提高内核代码的性能。通过理解并应用这些调度技术，开发者可以更好地优化其计算任务。
摘要由CSDN通过智能技术生成

文章目录

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中的调度原语，它允许用户轻松灵活地调度计算。

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

  • 通过一系列操作描述您的计算。
  • 尝试使用原语安排计算。
  • 编译并运行以查看性能差异。
  • 根据运行结果调整你的调度。
Mars_WH
关注
  • 6
    点赞
  • 踩
  • 25
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
专栏目录
TVM学习教程
qq_33287871的博客
01-29 2676
目录Step 1：TVM 视频介绍Step 2：研读 TVM 论文Step 3：了解 TVM 官网中可以学习的资料Step 4：上手 TVM CodeStep 5：TVM 其他相关资料 Step 1：TVM 视频介绍 https://www.bilibili.com/video/av29521815/ Step 2：研读 TVM 论文 https://www.usenix.org/system/files/osdi18-chen.pdf Step 3：了解 TVM 官网中可以学习的资料 https://tv
【TVM系列二】TVM介绍
Flying Cat
11-23 5572
文章同步更新在公众号 AIPlayer，欢迎扫码关注，共同进步 目录 一、TVM的工作流程 1、整体流程 2、关键数据结构 3、Transformations 4、搜索空间和基于机器学习的转换 5、目标代码转化 二、逻辑架构组件 三、运行TVM实例 1、交叉编译runtime 2、编译模型 3、运行模型 四、总结 一、TVM的工作流程 TVM主要由两个部分组成： （1）TVM编译器：负责编译和优化模型 （2）TVM runtime：提供目标设备上运行模型的API ..
TVM中的调度原语
吴建明wujianming_110117
12-14 634
TVM中的调度原语 TVM是一种用于高效内核构造的领域专用语言。 本文将展示如何通过TVM提供的各种原语来调度计算。 from future import absolute_import, print_function import tvm from tvm import te import numpy as np 通常存在多种方法来计算相同的结果，但是不同的方法会导致不同的局部性和性能。因此TVM要求用户提供如何执行称为Schedule的计算。 调度是一组计算转换，它转换程序中的计算循环。 declare
TVM Compiler中文教程：使用TVM编写可调模板和使用自动调优器
Mars_WH的专栏
06-14 4883
文章目录使用TVM编写可调模板和使用自动调优器加载依赖库第一步：定义搜索空间参数化调度使用更好的空间定义API第二步：搜索空间在TVM中使用自动调优器开始调优 使用TVM编写可调模板和使用自动调优器 这是TVM中auto-tuning模块的入门教程。 auto-tuning分两个步骤：第一步定义搜索空间；第二步是运行搜索算法来探索这个空间。在本教程中，你可以了解如何在TVM中执行这两个步骤。下面通...
tvm编译器原理与实践 实战踩坑
最新发布
离开那天的博客
06-13 331
首先要cp cmake的config.cmake到build文件夹下。首先下载resnet50-v2-7.onnx。然后我修改了.gitsubmodule文件。这一步可以替换成gitee的网址。然后.bashrc上面加上这一句。这样获得了resnet50.so。我要安装v0.14.0的版本。recursive是很重要的。今天玩了一下tvm的安装。这些第三方库下载完毕之后。make -j4开始构建。参考下面的method1。这时候你会发现有一些日志。之后直接cmake…如果你不喜欢这些日志。
TVM Compiler中文教程：TVM为NVIDIA GPU自动调优卷积网络
Mars_WH的专栏
06-16 3345
文章目录TVM为NVIDIA GPU自动调优卷积网络加载依赖库定义网络设置调优选项开始调优使用多个设备扩展测量范围 TVM为NVIDIA GPU自动调优卷积网络 针对特定设备和工作负载进行自动调优对于获得最佳性能至关重要。这篇教程是关于TVM如何为NVIDIA GPU调优整个卷积网络。 TVM中NVIDIA GPU的算子实现是以模板template形式编写的。模板有许多可调旋钮knobs（平铺ti...
TVM自动调度器
吴建明wujianming_110117
03-12 618
TVM自动调度器 随着模型大小，算子多样性和硬件异构性的不断增长，优化深度神经网络的执行速度非常困难。从计算的角度来看，深度神经网络只是张量计算的一层又一层。这些张量计算（例如matmul和conv2d）可以通过数学表达式轻松描述。在现代硬件上为其提供高性能的实现可能会非常具有挑战性。必须应用各种低级优化，利用特殊的硬件内在函数来实现高性能。建立线性代数和神经网络加速库（如CuBLAS，CuDNN，oneMKL和oneDNN）需要大量的工程工作。 如果可以编写数学表达式，将其神奇地转化为有效的代码实现，那么
全文翻译（二）: TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
吴建明wujianming_110117
10-27 516
全文翻译（二）: TVM: An Automated End-to-End Optimizing Compiler for Deep Learning 3．优化计算图 计算图是在DL框架中表示程序的常用方法。图3显示了两层卷积神经网络的计算图表示示例。这种高级表示与低级编译器中间表示（IR，如LLVM）间的主要区别，中间数据项是大型多维张量。计算图提供了算子的全局视图，但避免指定每个算子必须如何实现。与LLVM IRs一样，计算图可以转换为功能等效图优化。利用常见DL工作负载中的形状特异性，优化一组固定的输
全文翻译（全文合集）：TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
吴建明wujianming_110117
10-29 741
全文翻译（全文合集）：TVM: An Automated End-to-End Optimizing Compiler for Deep Learning 摘要 人们越来越需要将机器学习应用到各种各样的硬件设备中。现在的框架依赖于特定于供应商的算子库，针对窄带的服务器级GPU进行优化。将工作负荷部署到新平台，如移动电话，嵌入式设备和加速算子（如FPGA，ASIC）-需要大量手动操作。 提出了TVM，一个开源图形级别的编译器和算子级优化，提供高效跨多个应用程序移植到深度学习工作负荷硬件后端。TVM解决了深度学
全文翻译（一）：TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
吴建明wujianming_110117
10-26 529
全文翻译（一）：TVM: An Automated End-to-End Optimizing Compiler for Deep Learning 摘要 人们越来越需要将机器学习应用到各种各样的硬件设备中。现在的框架依赖于特定于供应商的算子库，针对窄带的服务器级GPU进行优化。将工作负荷部署到新平台，如移动电话，嵌入式设备和加速算子（如FPGA，ASIC）-需要大量手动操作。 提出了TVM，一个开源图形级别的编译器和算子级优化，提供高效跨多个应用程序移植到深度学习工作负荷硬件后端。TVM解决了深度学习特有
TVM: End-to-End Optimization Stack for Deep Learning
诗蕊的专栏
08-22 3227
Optimizing Computational Graphs Computational Graph：计算图，计算图用于描绘一个整体的计算任务，会根据图中的计算任务提前分配内存，所以在计算图中做优化，即将原来的计算图转换为相同效果的计算图。例如有些计算可以提前进行然后存储在内存中在需要用的时候直接调用，这样就可以节省一部分执行的时间。下图描述了一些计算图优化的例子： Operator F...
深度学习推理框架tvm-RK3588-gpu使用完整代码
12-22
1. 嵌入式设备RK3588-GPU使用 2. 深度学习框架-tvm 3. tvm+rk3588-gpu加速模型的推理 3.1 onnx模型的转换 3.2 tvm调用转换的模型
【tvm官网教程】张量表达与调度
B10030607的博客
04-21 1938
【tvm官网教程】张量表达与调度目的1. 调度原语1.1 te常用接口1.2 stage常用成员函数欢迎使用Markdown编辑器新的改变功能快捷键合理的创建标题，有助于目录的生成如何改变文本的样式插入链接与图片如何插入一段漂亮的代码片生成一个适合你的列表创建一个表格设定内容居中、居左、居右SmartyPants创建一个自定义列表如何创建一个注脚注释也是必不可少的KaTeX数学公式新的甘特图功能，丰富你的文章UML 图表FLowchart流程图导出与导入导出导入 目的 编译器是下一代大厂之争的关键，曾经有过
TVM之Schedule
sayandroid的专栏
07-07 4049
TVM之Schedule1. schedule2. stage3. methods for stage3.1 split3.2 tile3.3 fuse3.4 reorder3.5 bind3.6 compute_at3.7 compute_inline3.8 compute_root 最近在阅读TVM官网的文档，在阅读的过程中将keywords以及相应理解记录了下来，有理解不对的地方，欢迎批评指...
TVM Compiler中文教程：TVM中Compute和Reduction如何使用元组输入
Mars_WH的专栏
06-07 1691
Compute和Reduction使用元组输入 我们通常希望在单个循环内计算具有相同维度的多个输出，或者，执行涉及argmax等多个值的缩减。 在这篇教程，我们将介绍在TVM中元组输入。 from __future__ import absolute_import, print_function import tvm import numpy as np batch计算 对于具有相同维度的运算...
compute和schedule通过scheduler_client进行rpc通讯
jason的笔记
11-16 734
compute和schedule 节点之前都是通过rpc来通讯的，例如在E:\nova\nova\compute\manager.py 中一般通过self.scheduler_client 来和schedule 通讯，例如下面的例子 def _update_scheduler_instance_info(self, context, instance): """Sends
性能优化10_JobScheduler源码
baopengjian的专栏
01-26 308
1 JobScheduler 实现获取 JobScheduler tm = (JobScheduler) getSystemService(Context.JOB_SCHEDULER_SERVICE); 两个方面分析： 系统服务什么时候启动的 寻找是否有JobSchedulerService的类 结果：发现JobSchedulerService 存在 JobSchedulerService...
TVM代码库演练示例
weixin_42164269的博客
02-13 1343
TVM代码库演练示例 了解新的代码库可能是一个挑战。对于像TVM这样的代码库，尤其如此，其中不同的组件以非显而易见的方式交互。在本指南中，我们尝试通过一个简单的示例来说明构成编译管道的关键元素。对于每个重要步骤，我们都会显示在代码库中的哪个位置。目的是让新开发人员和感兴趣的用户更快地进入代码库。 代码库结构概述 在TVM存储库的根目录中，我们具有以下子目录，这些子目录一起构成了大部分代码库...
TVM节点属性更新方案：C++暴露方法和Python绑定
TVM 节点属性更新方案 TVM（TVM 是 Apache 的开源机器学习编译器堆栈）是一种机器学习编译器，它可以将机器学习模型转换为高效的机器代码。为了实现机器学习模型的优化和部署，TVM 提供了节点属性更新机制，该机制...
评论 1
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值