TVM Compiler中文教程：TVM使用张量化Tensorize利用硬件内联函数

TVM使用Tensorize利用硬件内联函数

张量化计算优化这篇教程没看太懂。
这篇教程是关于在TVM中如何执行张量化的介绍。

通过使用调度原语tensorize，人们可以用相应的内联函数替换计算单元，从而可以轻松利用手工制作的微内核函数，和扩展TVM来支持新的硬件架构。

本教程的目的是展示tensorize的功能和用法，而不是提供有效的解决方案。

from __future__ import absolute_import, print_function

import tvm
import numpy as np

定义矩阵乘法

以矩阵乘法为例。 矩阵乘法首先将两个矩阵之间的相应元素相乘，然后在某个轴上累加。下面，在TVM中描述了A*B^T的计算。

N, M, L = 1024, 512, 64
A = tvm.placeholder((N, L), name='A')
B = tvm.placeholder((M, L), name='B')
#设置reduce轴
k = tvm.reduce_axis((0, L), name='k')
C = tvm.compute((N, M), lambda i, j:
                tvm.sum(A[i, k] * B[j, k], axis=k), name='C')
s = tvm.create_schedule(C.op)
print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C {
  for (i, 0, 1024) {
    for (j, 0, 512) {
      C[((i*512) + j)] = 0.000000f
      for (k, 0, 64) {
        C[((i*512) + j)] = (C[((i*512) + j)] + (A[((i*64) + k)]*B[((j*64) + k)]))
      }
    }
  }
}

调度矩阵乘法

现在，假设我们有一个加速器支持矩阵向量乘法（GEMV）作为硬件原语，它可以采用任意大小的reduce轴，但另一个轴需要不大于16.因此我们分解矩阵乘法循环使最里面的循环为(16x64)GEMV。

factor = 16
# x,y对应N,M
x, y = C.op.axis
#z对应L,sum缩减
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))

输出：

produce C {
  for (i, 0, 1024) {
    for (j.outer, 0, 32) {
      for (j.inner, 0, 16) {
        C[((((i*32) + j.outer)*16) + j.inner)] = 0.000000f
        for (k, 0, 64) {
          C[((((i*32) + j.outer)*16) + j.inner)] = (C[((((i*32) + j.outer)*16) + j.inner)] + (A[((i*64) + k)]*B[((((j.outer*16) + j.inner)*64) + k)]))
        }
      }
    }
  }
}

如上IP打印显示，内循环j.inner与k和起来形成一个GEMV计算。在最内层的两个循环内，索引i是固定的，对矩阵A的访问仅改变k，这使得A的访问模式成为“向量”。为了利用我们假设的硬件GEMV指令，我们可以对j.inner进行张量化。

定义GEMV张量内联函数

在调度张量化之前，我们需要首先定义GEMV的内联函数。它包括两部分，第一部分是GEMV的计算定义，TVM使用它来匹配原始矩阵乘法调度中的计算模式；第二个是指定如何在设备上执行GEMV，这在下面的intrin_func中完成。

def intrin_gemv(m, l):
    a = tvm.placeholder((l,), name='a')
    b = tvm.placeholder((m,l), name='b')
    k = tvm.reduce_axis((0,l), name='k')
    c = tvm.compute((m,), lambda i: tvm.sum(a[k]*b[i,k],axis=k), name='c')
    #声明buffer
    Ab = tvm.decl_buffer(a.shape, a.dtype,
                        name='A',
                        offset_factor=1,
                        stride=[1])
    Bb = tvm.decl_buffer(b.shape, b.dtype,
                         name="B",
                         offset_factor=1,
                         strides=[tvm.var("s1"), 1])
    Cb = tvm.decl_buffer(c.shape, c.dtype,
                         name="C",
                         offset_factor=1,
                         strides=[1])
    #设备原生函数调用
    def intrin_func(ins,outs):
        ib = tvm.ir_builder.create()
        aa, bb = ins
        cc = outs[0]
        ib.emit(tvm.call_extern("int32","gemv_update",
               cc.access_ptr('w'),
               aa.access_ptr('r'),
               bb.access_ptr('r'),
               m,l,bb.stride[0]))
        return ib.get()
    #
    with tvm.build_config(offset_factor=1):
        return tvm.decl_tensor_intrin(c.op, intrin_func, binds={a: Ab, b: Bb, c: Cb})
        

这里tvm.decl_tensor_intrin声明怎么去执行C.op计算。我们的实现只需简单的输入和输出，将它们转换为指针并发出外部硬件内联函数调用。请注意，张量化需要用户指定offset_factor，利用此信息，TVM知道在原始数据结构的起始地址和传递给张量的偏移之间的数据是否被对齐，因此它有机会使用矢量化加载进行优化。为简化起见，我们将因子设置为1。

用于输入和输出的缓冲区也被声明，但这不是必需的，我们受益于缓冲区提供的额外信息。例如，我们将bb.strides [0]作为参数传递给外部函数gemv_update。现在bb.strides [0] == l，但稍后我们将看到它们如何与更复杂的调度有所不同。

注意，我们使用tvm.var(“s1”)作为B的第一个步幅维度。如果可以推断出步幅 - 在这种情况下，TVM知道张量B是紧凑的，因此步幅是[L，1] - 这样的占位符可以让TVM自动绑定我们的推断值。

#gemv调用
gemv = intrin_gemv(factor, L)
#张量化
s[C].tensorize(yi, gemv)
print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C {
  for (i, 0, 1024) {
    for (j.outer, 0, 32) {
      gemv_update(tvm_address_of(C[(((i*32) + j.outer)*16)]), tvm_address_of(A[(i*64)]), tvm_address_of(B[(j.outer*1024)]), 16, 64, 64)
    }
  }
}

通过对yi进行张量化，最内部的两个循环现在被我们之前定义的内联函数所取代。为了构建和运行模块，让我们定义外部函数gemv_update，它是一个简单的GEMV实现，仅用于演示。

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 util, clang
    temp = util.tempdir()
    ll_path = temp.relpath("temp.ll")
    # 从C源码创建LLVM IR
    ll_code = clang.create_llvm(cc_code, output=ll_path)
    return ll_code

现在我们利用编译属性import_llvm来导入llvm内联汇编。导入需要在张量化的GEMV执行之前进行。

s[C].pragma(x, "import_llvm", gemv_impl())
print(tvm.lower(s, [A, B, C], simple_mode=True))

输出：

produce C {
  // attr [iter_var(i, )] pragma_import_llvm = "; ModuleID = '/tmp/tmps1q05tf9/input0.cc'\nsource_filename = \"/tmp/tmps1q05tf9/input0.cc\"\ntarget datalayout = \"e-m:e-i64:64-f80:128-n8:16:32:64-S128\"\ntarget triple = \"x86_64-pc-linux-gnu\"\n\n; Function Attrs: noinline nounwind optnone uwtable\ndefine i32 @gemv_update(float*, float*, float*, i32, i32, i32) #0 {\n  %7 = alloca float*, align 8\n  %8 = alloca float*, align 8\n  %9 = alloca float*, align 8\n  %10 = alloca i32, align 4\n  %11 = alloca i32, align 4\n  %12 = alloca i32, align 4\n  %13 = alloca i32, align 4\n  %14 = alloca i32, align 4\n  store float* %0, float** %7, align 8\n  store float* %1, float** %8, align 8\n  store float* %2, float** %9, align 8\n  store i32 %3, i32* %10, align 4\n  store i32 %4, i32* %11, align 4\n  store i32 %5, i32* %12, align 4\n  store i32 0, i32* %13, align 4\n  br label %15\n\n; <label>:15:                                     ; preds = %50, %6\n  %16 = load i32, i32* %13, align 4\n  %17 = load i32, i32* %10, align 4\n  %18 = icmp slt i32 %16, %17\n  br i1 %18, label %19, label %53\n\n; <label>:19:                                     ; preds = %15\n  store i32 0, i32* %14, align 4\n  br label %20\n\n; <label>:20:                                     ; preds = %46, %19\n  %21 = load i32, i32* %14, align 4\n  %22 = load i32, i32* %11, align 4\n  %23 = icmp slt i32 %21, %22\n  br i1 %23, label %24, label %49\n\n; <label>:24:                                     ; preds = %20\n  %25 = load float*, float** %8, align 8\n  %26 = load i32, i32* %14, align 4\n  %27 = sext i32 %26 to i64\n  %28 = getelementptr inbounds float, float* %25, i64 %27\n  %29 = load float, float* %28, align 4\n  %30 = load float*, float** %9, align 8\n  %31 = load i32, i32* %13, align 4\n  %32 = load i32, i32* %12, align 4\n  %33 = mul nsw i32 %31, %32\n  %34 = load i32, i32* %14, align 4\n  %35 = add nsw i32 %33, %34\n  %36 = sext i32 %35 to i64\n  %37 = getelementptr inbounds float, float* %30, i64 %36\n  %38 = load float, float* %37, align 4\n  %39 = fmul float %29, %38\n  %40 = load float*, float** %7, align 8\n  %41 = load i32, i32* %13, align 4\n  %42 = sext i32 %41 to i64\n  %43 = getelementptr inbounds float, float* %40, i64 %42\n  %44 = load float, float* %43, align 4\n  %45 = fadd float %44, %39\n  store float %45, float* %43, align 4\n  br label %46\n\n; <label>:46:                                     ; preds = %24\n  %47 = load i32, i32* %14, align 4\n  %48 = add nsw i32 %47, 1\n  store i32 %48, i32* %14, align 4\n  br label %20\n\n; <label>:49:                                     ; preds = %20\n  br label %50\n\n; <label>:50:                                     ; preds = %49\n  %51 = load i32, i32* %13, align 4\n  %52 = add nsw i32 %51, 1\n  store i32 %52, i32* %13, align 4\n  br label %15\n\n; <label>:53:                                     ; preds = %15\n  ret i32 0\n}\n\nattributes #0 = { noinline nounwind optnone uwtable \"correctly-rounded-divide-sqrt-fp-math\"=\"false\" \"disable-tail-calls\"=\"false\" \"less-precise-fpmad\"=\"false\" \"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\"=\"+fxsr,+mmx,+sse,+sse2,+x87\" \"unsafe-fp-math\"=\"false\" \"use-soft-float\"=\"false\" }\n\n!llvm.module.flags = !{!0}\n!llvm.ident = !{!1}\n\n!0 = !{i32 1, !\"wchar_size\", i32 4}\n!1 = !{!\"clang version 6.0.1-svn334776-1~exp1~20190309042730.123 (branches/release_60)\"}\n"
  for (i, 0, 1024) {
    for (j.outer, 0, 32) {
      gemv_update(tvm_address_of(C[(((i*32) + j.outer)*16)]), tvm_address_of(A[(i*64)]), tvm_address_of(B[(j.outer*1024)]), 16, 64, 64)
    }
  }
}

最后，我们将张量化版本与numpy.dot生成的版本进行比较，确保我们的实现是正确的。

func = tvm.build(s, [A, B, C], target="llvm", name="gemv")

from topi.util 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)

为张量化进行Reduce-update

到目前为止，您已经学会了张量化的基本概念，现在让我们向更复杂的情况迈进一步。

假设我们的加速器只能将向量乘以一个方阵，其中向量大小不需要大于16.鉴于这样的硬件约束，现在我们需要将reduce轴分割如下，

zo, zi = s[C].split(z, factor=factor)
s[C].reorder(x, yo, zo, yi, zi)

但是，由于张量化内联计算现在只覆盖reduce轴的一部分，而不是使用一个“body”函数，TVM将在reduce循环之前需要调用一个reduce_reset函数，而reduce_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 util, clang
    temp = util.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 = tvm.placeholder((l,), name='a')
    b = tvm.placeholder((m, l), name='b')
    k = tvm.reduce_axis((0, l), name='k')
    c = tvm.compute((m,), lambda i:
    tvm.sum(a[k] * b[i, k], axis=k), name='c')
    Ab = tvm.decl_buffer(a.shape, a.dtype,
                         name="A",
                         offset_factor=1,
                         strides=[1])
    Bb = tvm.decl_buffer(b.shape, b.dtype,
                         name="B",
                         offset_factor=1,
                         strides=[tvm.var("s1"), 1])
    Cb = tvm.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.ir_builder.create()
            ib.emit(tvm.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.ir_builder.create()
            ib.emit(tvm.call_extern("int32", "gemv_reset", cc.access_ptr("w"), m))
            return ib.get()
        def _reduce_update():
            return _body()
        #返回三个函数
        return _body(), _reduce_reset(), _reduce_update()
        with tvm.build_config(offset_factor=1):
        return tvm.decl_tensor_intrin(c.op, intrin_func, binds={a: Ab, b: Bb, c: Cb})

请注意，现在intrin_func函数返回一个三元组:( body，reduce_reset，reduce_update）。如果张量化计算包含所有reduce轴，则将调用函数body()，否则将使用reduce_reset()和reduce_update()。在我们的示例中，body()和reduce_update()共享相同的实现，而在其他情况下，硬件可能对这两个函数有不同的指令。此外，我们可以看到，由于平铺，bb.strides [0]与l是不同的。

对方形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为用户提供了一种通过微内核获得完全优化调度的方法。例如，Intel CPU上的INT8量化可以使用张量来直接调用AVX指令。它还使TVM能够编译到ASIC-详情查看VTA。我们还演示了如何使用内联汇编导入，这有助于用户将汇编轻松导入到计算调度中。