使用Tensorize评估硬件内部特性

这是有关如何在TVM中执行张量的入门文档。

通过使用调度原语tensorize,人们可以用相应的内部函数代替计算单元,从而轻松利用handcrafted micro-kernels,扩展TVM以支持新的硬件体系结构。

本文的目的是展示张量的功能和用法,而不是提供有效的解决方案。

from __future__ import absolute_import, print_function
 
import tvm
from tvm import te
import numpy as np

定义矩阵乘法

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

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

输出:

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
             B: Buffer(B_2: Pointer(float32), float32, [512, 64], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 64], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i: int32, 0, 1024) {
    for (j: int32, 0, 512) {
      C_2[((i*512) + j)] = 0f32
      for (k: int32, 0, 64) {
        C_2[((i*512) + j)] = ((float32*)C_2[((i*512) + j)] + ((float32*)A_2[((i*64) + k)]*(float32*)B_2[((j*64) + k)]))
      }
    }
  }
}

调度Matmul

假设有一个支持矩阵矢量乘法(GEMV)作为硬件原语的加速器,可以采用任意大小的reduce轴,但另一个轴必须不大于16。因此,分解了matmul循环,生成最里面的一个(16x64)GEMV循环。

factor = 16
x, y = C.op.axis
(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))

输出:

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
             B: Buffer(B_2: Pointer(float32), float32, [512, 64], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 64], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i: int32, 0, 1024) {
    for (j.outer: int32, 0, 32) {
      for (j.inner: int32, 0, 16) {
        C_2[(((i*512) + (j.outer*16)) + j.inner)] = 0f32
        for (k: int32, 0, 64) {
          C_2[(((i*512) + (j.outer*16)) + j.inner)] = ((float32*)C_2[(((i*512) + (j.outer*16)) + j.inner)] + ((float32*)A_2[((i*64) + k)]*(float32*)B_2[(((j.outer*1024) + (j.inner*64)) + k)]))
        }
      }
    }
  }
}

如上面打印的IR所示,内部循环j.inner与k一起形成GEMV的计算-在最内部的两个循环内,索引i是固定的,对矩阵的访问A仅变化k,生成A“向量”的访问模式”。 可以用j.inner张量来评估假设的硬件的GEMV指令。

定义固有的GEMV张量化

调度张量前,先定义GEMV的固有函数。它包括两部分,第一部分是GEMV的计算定义。TVM使用它来匹配原始Matmul调度中的计算模式。第二个是指定如何在设备上执行GEMV,这在intrin_func下面完成。

def intrin_gemv(m, l):
    a = te.placeholder((l,), name="a")
    b = te.placeholder((m, l), name="b")
    k = te.reduce_axis((0, l), name="k")
    c = te.compute((m,), lambda i: te.sum(a[k] * b[i, k], axis=k), name="c")
    Ab = tvm.tir.decl_buffer(a.shape, a.dtype, name="A", offset_factor=1, strides=[1])
    Bb = tvm.tir.decl_buffer(b.shape, b.dtype, name="B", offset_factor=1, strides=[te.var("s1"), 1])
    Cb = tvm.tir.decl_buffer(c.shape, c.dtype, name="C", offset_factor=1, strides=[1])
 
    def intrin_func(ins, outs):
        ib = tvm.tir.ir_builder.create()
        aa, bb = ins
        cc = outs[0]
        ib.emit(
            tvm.tir.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()
 
    return te.decl_tensor_intrin(c.op, intrin_func, binds={a: Ab, b: Bb, c: Cb})

在此te.decl_tensor_intrin声明如何执行计算c.op。实现只接受输入和输出,将它们转换为指针并发出外部函数调用。注意,张量需要用户指定offset_factor,原始数据结构的起始地址和传递给张量的偏移量之间对齐的问题,TVM能评估,通过矢量化加载进行优化。为了简化,将系数设置为1。

为输入和输出声明了缓冲区,尽管这不是必需的,将从缓冲区提供的额外信息中受益。例如,bb.strides[0]作为参数传递 给外部函数gemv_update。将看到bb.strides[0] == l如何与更复杂的调度区分开。

注意,将te.var("s1")用作第一个步幅B。如果可以推理出步幅(在这种情况下,TVM确定张量B是紧凑的,步幅是[L, 1]),可以使用此类placeholder让TVM自动为绑定推理的值。

gemv = intrin_gemv(factor, L)
s[C].tensorize(yi, gemv)
print(tvm.lower(s, [A, B, C], simple_mode=True))

输出:

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
             B: Buffer(B_2: Pointer(float32), float32, [512, 64], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 64], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  for (i: int32, 0, 1024) {
    for (j.outer: int32, 0, 32) {
      @tir.call_extern("gemv_update", @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), C_2, ((i*512) + (j.outer*16)), 16, 2, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), A_2, (i*64), 64, 1, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), B_2, (j.outer*1024), 1024, 1, dtype=handle), 16, 64, 64, dtype=int32)
    }
  }
}

通过张大yi,最里面的两个循环被之前定义的内在函数代替。为了构建和运行该模块,定义外部函数gemv_update,它是GEMV的naive实现,仅用于演示。

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 utils, clang
 
    temp = utils.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

利用pragma属性import_llvm导入llvm asm内联。导入在执行张量的GEMV之前进行。

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

出:

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 512], []),
             B: Buffer(B_2: Pointer(float32), float32, [512, 64], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 64], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  attr [IterVar(i: int32, (nullptr), "DataPar", "")] "pragma_import_llvm" = "; ModuleID = '/tmp/tmpinr5hwkd/input0.cc'
source_filename = "/tmp/tmpinr5hwkd/input0.cc"
target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
target triple = "x86_64-pc-linux-gnu"
 
; Function Attrs: noinline nounwind optnone uwtable
define dso_local i32 @gemv_update(float*, float*, float*, i32, i32, i32) #0 {
  %7 = alloca float*, align 8
  %8 = alloca float*, align 8
  %9 = alloca float*, align 8
  %10 = alloca i32, align 4
  %11 = alloca i32, align 4
  %12 = alloca i32, align 4
  %13 = alloca i32, align 4
  %14 = alloca i32, align 4
  store float* %0, float** %7, align 8
  store float* %1, float** %8, align 8
  store float* %2, float** %9, align 8
  store i32 %3, i32* %10, align 4
  store i32 %4, i32* %11, align 4
  store i32 %5, i32* %12, align 4
  store i32 0, i32* %13, align 4
  br label %15
 
15:                                               ; preds = %50, %6
  %16 = load i32, i32* %13, align 4
  %17 = load i32, i32* %10, align 4
  %18 = icmp slt i32 %16, %17
  br i1 %18, label %19, label %53
 
19:                                               ; preds = %15
  store i32 0, i32* %14, align 4
  br label %20
 
20:                                               ; preds = %46, %19
  %21 = load i32, i32* %14, align 4
  %22 = load i32, i32* %11, align 4
  %23 = icmp slt i32 %21, %22
  br i1 %23, label %24, label %49
 
24:                                               ; preds = %20
  %25 = load float*, float** %8, align 8
  %26 = load i32, i32* %14, align 4
  %27 = sext i32 %26 to i64
  %28 = getelementptr inbounds float, float* %25, i64 %27
  %29 = load float, float* %28, align 4
  %30 = load float*, float** %9, align 8
  %31 = load i32, i32* %13, align 4
  %32 = load i32, i32* %12, align 4
  %33 = mul nsw i32 %31, %32
  %34 = load i32, i32* %14, align 4
  %35 = add nsw i32 %33, %34
  %36 = sext i32 %35 to i64
  %37 = getelementptr inbounds float, float* %30, i64 %36
  %38 = load float, float* %37, align 4
  %39 = fmul float %29, %38
  %40 = load float*, float** %7, align 8
  %41 = load i32, i32* %13, align 4
  %42 = sext i32 %41 to i64
  %43 = getelementptr inbounds float, float* %40, i64 %42
  %44 = load float, float* %43, align 4
  %45 = fadd float %44, %39
  store float %45, float* %43, align 4
  br label %46
 
46:                                               ; preds = %24
  %47 = load i32, i32* %14, align 4
  %48 = add nsw i32 %47, 1
  store i32 %48, i32* %14, align 4
  br label %20
 
49:                                               ; preds = %20
  br label %50
 
50:                                               ; preds = %49
  %51 = load i32, i32* %13, align 4
  %52 = add nsw i32 %51, 1
  store i32 %52, i32* %13, align 4
  br label %15
 
53:                                               ; preds = %15
  ret i32 0
}
 
attributes #0 = { noinline nounwind optnone uwtable "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "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"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" }
 
!llvm.module.flags = !{!0}
!llvm.ident = !{!1}
 
!0 = !{i32 1, !"wchar_size", i32 4}
!1 = !{!"clang version 9.0.1-+20191211110317+c1a0a213378-1~exp1~20191211221711.104 "}
";
  for (i, 0, 1024) {
    for (j.outer: int32, 0, 32) {
      @tir.call_extern("gemv_update", @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), C_2, ((i*512) + (j.outer*16)), 16, 2, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), A_2, (i*64), 64, 1, dtype=handle), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), B_2, (j.outer*1024), 1024, 1, dtype=handle), 16, 64, 64, dtype=int32)
    }
  }
}

最后,将张量版本与numpy.dot产生的张量版本进行比较,确保实现正确。

func = tvm.build(s, [A, B, C], target="llvm", name="gemv")
 
from tvm.topi.utils 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)

进行Tensorize更新Reduce-update

已经了解了张量化的基本概念,现在让向更复杂的情况迈进一步。

假设加速器只能将向量乘以一个矩阵,向量的大小必须不大于16。考虑到硬件限制,需要按如下方式拆分reduce轴。

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

由于张量内在函数现在仅覆盖了reduce轴的一部分,而不是使用一个“ body”函数,因此TVM需要一个reduce_reset在reduce for循环之前调用的reduce_update函数,以及一个定义“ 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 utils, clang
 
    temp = utils.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 = te.placeholder((l,), name="a")
    b = te.placeholder((m, l), name="b")
    k = te.reduce_axis((0, l), name="k")
    c = te.compute((m,), lambda i: te.sum(a[k] * b[i, k], axis=k), name="c")
    Ab = tvm.tir.decl_buffer(a.shape, a.dtype, name="A", offset_factor=1, strides=[1])
    Bb = tvm.tir.decl_buffer(b.shape, b.dtype, name="B", offset_factor=1, strides=[te.var("s1"), 1])
    Cb = tvm.tir.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.tir.ir_builder.create()
            ib.emit(
                tvm.tir.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.tir.ir_builder.create()
            ib.emit(tvm.tir.call_extern("int32", "gemv_reset", cc.access_ptr("w"), m))
            return ib.get()
 
        def _reduce_update():
            return _body()
 
        return _body(), _reduce_reset(), _reduce_update()
 
    return te.decl_tensor_intrin(c.op, intrin_func, binds={a: Ab, b: Bb, c: Cb})

注意,intrin_func返回一个三元组: 如果tensorization包括所有的reduce轴,功能将被调用,否则一起将被使用。在示例中,共享相同的实现,而在其它情况下,硬件对于这两个功能可能具有不同的指令。此外,由于平铺,可以看到现在是不同的。(body, reduce_reset, reduce_update)body()reduce_reset()reduce_update()body()reduce_update()bb.strides[0]l

张量squared 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为用户提供了一种通过微内核获得完全优化的调度方式。例如,英特尔CPU上使用张量化直接调用AVX指令进行INT8量化。使TVM可以编译为ASIC-有关详细信息,请参阅VTA:深度学习加速器堆栈。演示了如何使用内联程序集导入,这可以帮助用户轻松地将asm输入调度中。

使用Tensorize评估硬件内部特性的更多相关文章

  1. js的内部特性--属性

    使用方法:通过调用Object.defineProperty(对象,"对象属性",{}进行的操作}) 当对一个对象的属性的属性类型中vlue设置为一个值时,则这个对象的这个属性的值 ...

  2. 【翻译】.NET 5中的性能改进

    [翻译].NET 5中的性能改进 在.NET Core之前的版本中,其实已经在博客中介绍了在该版本中发现的重大性能改进. 从.NET Core 2.0到.NET Core 2.1到.NET Core ...

  3. 【翻译】.NET 5 Preview7发布

    今天,发布了.NET 5.0 Preview7.这是倒数第二个预览版本(在转移到RC之前).此时,大多数功能应该已经非常接近完成了.Single file和ARM64 intrinsics是两个花费了 ...

  4. Linux转发性能评估与优化-转发瓶颈分析与解决方式(补遗)

    补遗 关于网络接收的软中断负载均衡,已经有了成熟的方案,可是该方案并不特别适合数据包转发,它对server的小包处理非常好.这就是RPS.我针对RPS做了一个patch.提升了其转发效率. 下面是我转 ...

  5. Linux转发性能评估与优化(转发瓶颈分析与解决方式)

    线速问题 非常多人对这个线速概念存在误解. 觉得所谓线速能力就是路由器/交换机就像一根网线一样. 而这,是不可能的.应该考虑到的一个概念就是延迟. 数据包进入路由器或者交换机,存在一个核心延迟操作,这 ...

  6. Flash硬件原理

    1.2.1. 什么是Flash Flash全名叫做Flash Memory,从名字就能看出,是种数据存储设备,存储设备有很多类,Flash属于非易失性存储设备(Non-volatile Memory ...

  7. 深入GPU硬件架构及运行机制

    目录 一.导言 1.1 为何要了解GPU? 1.2 内容要点 1.3 带着问题阅读 二.GPU概述 2.1 GPU是什么? 2.2 GPU历史 2.2.1 NV GPU发展史 2.2.2 NV GPU ...

  8. Microsoft Tech Summit 2018 课程简述:利用 Windows 新特性开发出更好的手绘视频应用

    概述 Microsoft Tech Summit 2018 微软技术暨生态大会将于10月24日至27日在上海世博中心举行,这也会是国内举办的最后一届 Tech Summit,2019 年开始会以 Mi ...

  9. CSS硬件加速的好与坏

    本文翻译自Ariya Hidayat的Hardware Accelerated CSS: The Nice vs The Naughty.感谢Kyle He帮助校对. 每个人都痴迷于60桢每秒的顺滑动 ...

随机推荐

  1. 请解释Spring事务传播传播行为

    什么叫事务传播行为? 听起来挺高端的,其实很简单. 即然是传播,那么至少有两个东西,才可以发生传播.单体不不存在传播这个行为. 事务传播行为(propagation behavior)指的就是当一个事 ...

  2. Linux中的防火墙(Netfilter、Iptables、Firewalld)

    目录 Netfilter Iptables iptables做本地端口转发 Firewalld Netfilter Netfilter是Linux 2.4内核引入的全新的包过滤引擎,位于Linux内核 ...

  3. Python脚本破解压缩文件口令(zipfile)

    环境:Windows python版本2.7.15 Python中操作zip压缩文件的模块是 zipfile . 相关文章:Python中zipfile压缩文件模块的使用 我们破解压缩文件的口令也是用 ...

  4. Jenkins 基础篇 - 环境准备

    前期准备 为了开发这套课程,我专门购买了一台二手戴尔服务器 Dell R720,同时把废弃多年的笔记本电脑也利用了起来.当然如果自己没有多余的电脑,也可以在自己电脑上安装 Vmware Worksta ...

  5. NumPy中文文档搬砖(划掉)学习笔记(1)

    原文地址 前言 况下加速Python中的操作运行时.适用于快速数值运算的一个选项是NumPy,它当之无愧地将自己称为使用Python进行科学计算的基本软件包. 当然,很少有人将50微秒(百万分之五十秒 ...

  6. 关于MySQL参数,这些你要知道

    前言: 在前面一些文章中,经常能看到介绍某某参数的作用,可能有些小伙伴仍搞不清楚 MySQL 参数是啥.本篇文章我们来聊聊 MySQL 参数,学习下如何管理维护 MySQL 参数. 1.MySQL参数 ...

  7. Kubernetes入门,使用minikube 搭建本地k8s 环境

    这是一篇 K8S 的 HelloWorld,在学习K8S官方文档时搭建环境搭建的一个记录,照着文档下来还是比较顺利的. 一.安装kubectl 下载 kubectl curl -LO "ht ...

  8. 项目展示$\alpha$

    项目 内容 课程:北航-2020-春-软件工程 博客园班级博客 要求 强制转会与项目展示 我们在这个课程的目标是 提升团队管理及合作能力,开发一项满意的工程项目 这个作业在哪个具体方面帮助我们实现目标 ...

  9. Envoy:主动健康监测

    实验文件 docker-compose version: '3' services: envoy: image: envoyproxy/envoy-alpine:v1.15-latest enviro ...

  10. [Java] 类库例题

    例1 字符串操作 定义一个StringBuffer类对象,然后通过append()方法向对象中添加26个小写字母,每次只添加一次,共添加26次,然后按逆序方式输出,并且可以删除前5个字符 面向过程实现 ...