首页 > 编程语言 >机器学习编译（三）：张量程序案例 TensorIR

机器学习编译（三）：张量程序案例 TensorIR

时间：2023-08-13 12:23:52浏览次数：42  
标签：vi vj vk 张量 编译 np 128 TensorIR float32

使用张量程序抽象的目的是为了表示循环和相关的硬件加速选择，如多线程、特殊硬件指令的使用和内存访问。

1. 一个例子

使用张量程序抽象，我们可以在较高层的抽象制定一些与特定硬件无关的较通用的 IR 优化（计算优化）。

比如，

对于两个大小为 128×128 的矩阵 A 和 B，我们进行如下两步的张量计算：

也就是做了一个 matmul 操作和一个 relu 操作。

使用 numpy 做这个操作，大致代码如下；

dtype = "float32"
a_np = np.random.rand(128, 128).astype(dtype)
b_np = np.random.rand(128, 128).astype(dtype)
# a @ b is equivalent to np.matmul(a, b)
c_mm_relu = np.maximum(a_np @ b_np, 0)

在底层，NumPy 调用库（例如 OpenBLAS）和它自己在低级 C 语言中的一些实现来执行这些计算。

类似的一个代码是这样：

def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
    Y = np.empty((128, 128), dtype="float32")
    for i in range(128):
        for j in range(128):
            for k in range(128):
                if k == 0:
                    Y[i, j] = 0
                Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
    for i in range(128):
        for j in range(128):
            C[i, j] = max(Y[i, j], 0)

这里相比原来的 numpy 代码做了内存分配和循环展开，比较接近底层语言做的事了。

验证一下两种实现的结果，发现是等价的：

c_np = np.empty((128, 128), dtype=dtype)
lnumpy_mm_relu(a_np, b_np, c_np)
np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)

但第二种实现更接近实际的框架在目标硬件上的调度和计算。这里的抽象展示了我们在实际实现中需要的元素：

  • 多维缓冲区（数组）。
  • 在数组维度上的循环。
  • 在循环下执行的计算语句。

2. TensorIR

接下来介绍 TensorIR，TensorIR 是 TVM 里的 Tensor 抽象，是最接近硬件的 IR。

在 TVM 中，可以用一种叫做 TVMScript DSL（Domain Specified Language，特定领域语言）来用 python 写一些 IR 优化，比如上面的例子可以写成：

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def mm_relu(A: T.Buffer((128, 128), "float32"),
                B: T.Buffer((128, 128), "float32"),
                C: T.Buffer((128, 128), "float32")):
        T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
        Y = T.alloc_buffer((128, 128), dtype="float32")
        for i, j, k in T.grid(128, 128, 128):
            with T.block("Y"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j)
                vk = T.axis.reduce(128, k)
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(128, 128):
            with T.block("C"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j)
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

对比一下用 numpy 指定的优化和用 TVMScript 的优化：

从这个程序出发我们看一下 TVMScript 的语法和他们对应的 Tensor 抽象：

  • 函数参数与缓冲区

    TVMScript 用 T.Buffer 指定数组大小：

    # TensorIR
def mm_relu(A: T.Buffer((128, 128), "float32"),
            B: T.Buffer((128, 128), "float32"),
            C: T.Buffer((128, 128), "float32")):
    ...
# numpy
def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
    ...

    这里 A、B、C 都是 128 × 128、元素类型为 float32 的 Tensor。制定数组大小有助于编译器对计算过程的调度和计算做更多的优化。

    计算过程中会产生中间结果 Y（matmul 之后，relu 之前），TVMScript 这里也对中间结果分配了缓冲区：

    # TensorIR
Y = T.alloc_buffer((128, 128), dtype="float32")
# numpy
Y = np.empty((128, 128), dtype="float32")

  • for 循环

    TVMScript 用 T.grid 来表示循环迭代：

    # TensorIR
for i, j, k in T.grid(128, 128, 128):
# numpy
for i in range(128):
    for j in range(128):
        for k in range(128):

  • 计算块

    CUDA 中也有 block 的概念，这里也是类似的，对于循环中的变量需要映射到一个块上做计算：

    # TensorIR
with T.block("Y"):
    vi = T.axis.spatial(128, i)
    vj = T.axis.spatial(128, j)
    vk = T.axis.reduce(128, k)
    with T.init():
        Y[vi, vj] = T.float32(0)
    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]

# coressponding numpy code
vi, vj, vk = i, j, k
if vk == 0:
    Y[vi, vj] = 0
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]

    这里的 matmul 映射到了一个块 Y 上得到计算结果，第 8 行给出了计算规则。

    块 是 TensorIR 中的基本计算单位。值得注意的是，该块包含比普通 NumPy 代码更多的信息。一个块包含一组块轴（vi、vj、vk）和围绕它们定义的计算。

    vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)

    上面三行声明了关于块轴的关键性质，语法如下。

    [block_axis] = T.axis.[axis_type]([axis_range], [mapped_value])

    这三行包含以下信息：

    • 定义了 vi、vj、vk 应被绑定到的位置（在本例中为 i、j 和 k）；
    • 声明了 vi、vj、vk 的原始范围（T.axis.spatial(128, i) 中的 128）；
    • 声明了块轴的属性（spatial, reduce）。

    我们一一浏览这些属性。首先，就边界关系而言，vi = T.axis.spatial(128, i) 有效地蕴含了 vi = i。[axis_range] 值提供了 [block_axis] 的预期范围。例如，vi = T.axis.spatial(128, i) 中的 128 表明 vi 应该在 range(0, 128) 中。

  • 块轴绑定的语法糖

    在每个块轴直接映射到外部循环迭代器的情况下，我们可以使用 T.axis.remap 在一行中声明所有块轴：

    # SSR means the properties of each axes are "spatial", "spatial", "reduce"
vi, vj, vk = T.axis.remap("SSR", [i, j, k])

    这里就相当于块 vi, vj, vk 分别和循环迭代器 i, j, k 绑定上了，类似 T.grid，只需要一行代码就可以完成多个绑定。

    使用这个语法之后的程序如下：

    @tvm.script.ir_module
class MyModuleWithAxisRemapSugar:
    @T.prim_func
    def mm_relu(A: T.Buffer((128, 128), "float32"),
                B: T.Buffer((128, 128), "float32"),
                C: T.Buffer((128, 128), "float32")):
        T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
        Y = T.alloc_buffer((128, 128), dtype="float32")
        for i, j, k in T.grid(128, 128, 128):
            with T.block("Y"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(128, 128):
            with T.block("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

  • 函数属性和装饰器

    @tvm.script.ir_module 表示 MyModule 是一个 IRModule。IRModule 是 TVM 中的最小编译单位（可以用 tvm.build 来编译）。

    一个 IRModule 中可以包含多个元张量函数 prim_func，比如这里 MyModuleWithTwoFunctions 有两个 prim_func：

    @tvm.script.ir_module
class MyModuleWithTwoFunctions:
    @T.prim_func
    def mm(A: T.Buffer((128, 128), "float32"),
           B: T.Buffer((128, 128), "float32"),
           Y: T.Buffer((128, 128), "float32")):
        T.func_attr({"global_symbol": "mm", "tir.noalias": True})
        for i, j, k in T.grid(128, 128, 128):
            with T.block("Y"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]

    @T.prim_func
    def relu(A: T.Buffer((128, 128), "float32"),
             B: T.Buffer((128, 128), "float32")):
        T.func_attr({"global_symbol": "relu", "tir.noalias": True})
        for i, j in T.grid(128, 128):
            with T.block("B"):
                vi, vj = T.axis.remap("SS", [i, j])
                B[vi, vj] = T.max(A[vi, vj], T.float32(0))

3. 变换

对于之前的 numpy 的 matmul + relu 的例子，在指定内存分配和循环展开的基础上还可以做一个等价的计算变换：

def lnumpy_mm_relu_v2(A: np.ndarray, B: np.ndarray, C: np.ndarray):
    Y = np.empty((128, 128), dtype="float32")
    for i in range(128):
        for j0 in range(32):
            for k in range(128):
                for j1 in range(4):
                    j = j0 * 4 + j1
                    if k == 0:
                        Y[i, j] = 0
                    Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
    for i in range(128):
        for j in range(128):
            C[i, j] = max(Y[i, j], 0)

c_np = np.empty((128, 128), dtype=dtype)
lnumpy_mm_relu_v2(a_np, b_np, c_np)
np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)

这里对于第二层循环 0 ~ 127 做了变换，用两个循环(32× 4) j0 和 j1 替换了 j 循环；

TVMScript 中提供了 Schedule 类可以帮助我们做类似的循环变换。

这是原先的未做循环变量的 TVM 程序：

from tvm.script import ir as I
from tvm.script import tir as T

@I.ir_module
class Module:
    @T.prim_func
    def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
        T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
        # with T.block("root"):
        Y = T.alloc_buffer((128, 128))
        for i, j, k in T.grid(128, 128, 128):
            with T.block("Y"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                T.reads(A[vi, vk], B[vk, vj])
                T.writes(Y[vi, vj])
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(128, 128):
            with T.block("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                T.reads(Y[vi, vj])
                T.writes(C[vi, vj])
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

我们首先创建一个以给定的 MyModule 作为输入的 Schedule 辅助类。

sch = tvm.tir.Schedule(MyModule)

然后我们执行以下操作以获得对块 Y 和相应循环的引用。

block_Y = sch.get_block("Y", func_name="mm_relu")
i, j, k = sch.get_loops(block_Y)

我们把对 j 的循环变换成两个循环，内部循环的长度是 4：

j0, j1 = sch.split(j, factors=[None, 4])

我们可以查看存储在 sch.mod 中的变换结果。

IPython.display.Code(sch.mod.script(), language="python")

变换后的 TVM 程序：

from tvm.script import ir as I
from tvm.script import tir as T

@I.ir_module
class Module:
    @T.prim_func
    def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
        T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
        # with T.block("root"):
        Y = T.alloc_buffer((128, 128))
        for i, j_0, j_1, k in T.grid(128, 32, 4, 128):
            with T.block("Y"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j_0 * 4 + j_1)
                vk = T.axis.reduce(128, k)
                T.reads(A[vi, vk], B[vk, vj])
                T.writes(Y[vi, vj])
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(128, 128):
            with T.block("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                T.reads(Y[vi, vj])
                T.writes(C[vi, vj])
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

变换之后，我们创建了两个新的循环，j_0 和 j_1，分别对应范围 32 和 4。下一步是重新排序这两个循环。

sch.reorder(j0, k, j1)
IPython.display.Code(sch.mod.script(), language="python")

重新排序之后的 TVM 程序：

from tvm.script import ir as I
from tvm.script import tir as T

@I.ir_module
class Module:
    @T.prim_func
    def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
        T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
        # with T.block("root"):
        Y = T.alloc_buffer((128, 128))
        for i, j_0, k, j_1 in T.grid(128, 32, 128, 4):
            with T.block("Y"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j_0 * 4 + j_1)
                vk = T.axis.reduce(128, k)
                T.reads(A[vi, vk], B[vk, vj])
                T.writes(Y[vi, vj])
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(128, 128):
            with T.block("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                T.reads(Y[vi, vj])
                T.writes(C[vi, vj])
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

可以发现这里 i, j_0, k, j_1 的迭代顺序和上面重新排序之前的 i, j_0, j_1, k 是不一样的。

这样，我们就用 TVMScript 对计算过程手动做了循环展开上的优化。

4. 编译和运行

前面我们说了 TensorIR 是 TVM 里最接近硬件的 IR，而 IRModule 是 TVM 中可以编译的最小单位。

也就是说，对于 MyModule，我们可以对它进行编译和运行。

TVM 中可以用 tvm.build 来编译指定的 IRModule，并且指定 target（目标硬件平台）：

rt_lib = tvm.build(MyModule, target="llvm")

我们创建三个用于保存输入和输出的 TVM NDArray：

a_nd = tvm.nd.array(a_np)
b_nd = tvm.nd.array(b_np)
c_nd = tvm.nd.empty((128, 128), dtype="float32")
type(c_nd)       # tvm.runtime.ndarray.NDArray

编译后的 rt_lib 可以运行，我们把输入扔进去：

func_mm_relu = rt_lib["mm_relu"]
func_mm_relu(a_nd, b_nd, c_nd)

np.testing.assert_allclose(c_mm_relu, c_nd.numpy(), rtol=1e-5)

可以发现运行结果和我们手写的 numpy 是一样的。

我们还可以构建经过变换后的程序：

rt_lib_after = tvm.build(sch.mod, target="llvm")
rt_lib_after["mm_relu"](a_nd, b_nd, c_nd)
np.testing.assert_allclose(c_mm_relu, c_nd.numpy(), rtol=1e-5)

最后，我们可以比较一下两者的时间差，看一下我们的变换的效果。 time_evaluator 是一个辅助的测试函数，可用于比较不同生成函数的运行性能：

f_timer_before = rt_lib.time_evaluator("mm_relu", tvm.cpu())
print("Time cost of MyModule %g sec" % f_timer_before(a_nd, b_nd, c_nd).mean)
f_timer_after = rt_lib_after.time_evaluator("mm_relu", tvm.cpu())
print("Time cost of transformed sch.mod %g sec" % f_timer_after(a_nd, b_nd, c_nd).mean)

Time cost of MyModule 0.00214587 sec
Time cost of transformed sch.mod 0.00101764 sec

看来我们的变换还是有很大的作用的，那么原因是什么呢？

如下图，现代 CPU 带有多级缓存，需要先将数据提取到缓存中，然后 CPU 才能访问它。

访问已经在缓存中的数据要快得多。CPU 采用的一种策略是获取彼此更接近的数据（局部性原理）。 当我们读取内存中的一个元素时，它会尝试将附近的元素（更正式的名称为“缓存行”）获取到缓存中。 因此，当你读取下一个元素时，它已经在缓存中。 因此，具有连续内存访问的代码通常比随机访问内存不同部分的代码更快。

现在让我们看看上面的迭代可视化，分析一下是怎么回事。 在这个分析中，让我们关注最里面的两个循环：k 和 j1。高亮的地方显示了当我们针对 k 的一个特定实例迭代 j1 时迭代触及的 Y、A 和 B 中的相应区域。

我们可以发现，j1 这一迭代产生了对 B 元素的连续访问。具体来说，它意味着在 j1=0 和 j1=1 时我们读取的值彼此相邻。这可以让我们拥有更好的缓存访问行为。此外，我们使 C 的计算更接近 Y，从而实现更好的缓存行为。

改变计算的循环迭代使之更接近 CPU 执行策略是编译器的常见优化手段之一。

除了 TensorIR，TVM 中还可以用张量表示（Tensor Expression, TE）来做张量的抽象，这里略过，在介绍 TVM 的文章中再详细介绍。

标签：vi,vj,vk,张量,编译,np,128,TensorIR,float32
From： https://www.cnblogs.com/linrj/p/17626385.html

相关文章

  • windows 10下编译mariadb
    环境:OS:Windows10mariadb:10.4.29 1.下载mariadb源码 2.安装依赖组件 3.修改sql_locale.cc文件编码文件位置:E:\mariadb\mariadb-10.4.29\sql在ue里面打开,另存为的时候将改成utf-8编码,选择格式为:utf-8withbom我这里只选择utf-8好像也可以 4.编译E:\mariadb>c......
  • JVM之字节码的编译原理
    JVM之字节码的编译原理Java最初诞生的目的就是为了在不依赖特定的物理硬件和操作系统环境下运行，那么也就是说Java程序实现跨平台我的基石其实就是字节码。Java之所以能够解决程序的安全性问题、跨平台移植性等问题，最主要的原因就是Java源代码的编译结果并非是本地机器指令，而是字......
  • 【我和openGauss的故事】构建openGauss开发编译提交一体化环境
    大数据模型[openGauss](javascript:void(0);)2023-07-2917:58发表于四川前文本文适合对openGauss源代码有好奇心的爱好者，那么一个友好的openGauss源代码环境应该是怎么样的。openGauss的开发环境是如何设置的？openGauss的编译环境是如何构建的？如何向openGauss提交代码，笔者集合官......
  • VS2019编译CloudCompare2.12.4
    参考：https://blog.csdn.net/pingfanderen5/article/details/1261800821.VisualStudio2019对应v142工具2.安装QT，qt5.14.2及以前的版本存在下载包，下载地址：http://download.qt.io/ ，但是5.14.2只能支持到VS2017。 所以采用在线安装的方式安装qt5.15.2版本 源码准......
  • 1、编译 glibc 过程中报错 ../configure --prefix=/opt/glibc-2.27       2、首
    64位安装包,查看系统位数,安装对应的mysqlLinux系统安装MySQL时，将MySQL-5.6.17-1.el6.x86_64.rpm-bundle.tar包打开，有7个rpm文件，如下：MySQL-client-5.6.17-1.el6.x86_64.rpmMySQL-devel-5.6.17-1.el6.x86_64.rpmMySQL-embedded-5.6.17-1.el6.x86_64.rpmMySQL-server-5.6.17-1.el6.......
  • openssl安装编译
    Ubuntuopenssl安装编译编译cmake时报错缺少openssl依赖[missing:OPENSSL_CRYPTO_LIBRARY]CMakeErrorat/usr/share/cmake-3.10/Modules/FindPackageHandleStandardArgs.cmake:137(message):CouldNOTfindOpenSSL,trytosetthepathtoOpenSSLrootfolderinthe......
  • 编译安装最新版本VIM
    编译安装最新版本VIM命令：gitclonehttps://github.com/vim/vim.gitcdvim/srcmakesudomakeinstallvim编译时可能出现的错误1.错误：noacceptableCcompilerfoundin$PATH解决办法:安装GCC软件套件atpinstallgcc2.编译vim时Youneedtoinstallaterminallib......
  • 如何在32位ubuntu11.10 下编译android 4.0.1源码和goldfish内核
    一准备工作 1安装javasdk6（1）从jdk官方网站http://www.oracle.com/technetwork/java/javase/downloads/jdk-6u29-download-513648.html下载jdk-6u29-linux-i586.bin文件。（2）执行jdk安装文件 [html] viewplaincopy1.$chmoda+xjdk-6u29-linux-i586.bin2.$jdk......
  • Ubuntu20.04 下编译和运行 FreeSWITCH的问题汇总
    1.Ubuntu20.04下编译和运行FreeSWITCH的问题汇总1.1.环境Ubuntu20.04.2LTS(Linux5.4.0-152-genericx86_64GNU/Linux)FreeSWITCH-1.10.9-release1.2.结论根据配置和编译过程中的错误提示，基本上就是一些依赖库的缺失问题，根据提示给出的依赖库及其版本要求，只要能在a......
  • python编译pyc文件
    python提供了内置的类库来实现把py文件编译为pyc文件，这个模块就是py_compile模块。将单个python文件转为pyc文件python-mpy_compilemycode.py将一个目录中的python文件转为pyc文件python-mcompileall./your_path/编译完成后如果想要直接运行Pyc文件注意两点：1.要把p......

赞助商

阅读排行

网站主页关于我们联系我们网站地图

本网站内容转载自其他媒体，侵权联系[admin##ips99.com]。

Copyright © 2020-2023 IPS99 版权所有 IPS99