使用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. hdu 1298 字典树 + DFS (模拟T9文本输入)

    题意:       给你一些按键顺序,让你输出每一步中概率最大的那个单词,这里的概率计算方 法好好看看别弄错了,一开始就是因为弄错了,各种wa,比如 abc 1 ,ab 1,那么 ab 的概率就是2 ...

  2. 关于HTTP的一些概念

    各种概念 HTTP HTTP(HyperText Transfer Protocol) -- 超文本传输协议 它可以拆成三个部分:"超文本"."传输".&quo ...

  3. 2021/5/11 docker的应用

    很快一天过去了,今天虽然没有加班,但是依旧感觉疲惫,弄了一天的docker镜像的问题.作为一名前端开发工程师从以前从未听说过docker到现在懂得了如何运用,也是不容易啊.之前也说过,我们项目结构是前 ...

  4. Jsp授课

    2.1 JSP基础 2.1.1 JSP简介 JSP全称是Java Server Page,是一种动态网页技术标准.它和Servlet一样,也是sun公司推出的一套开发动态web资源的技术,称为JSP/ ...

  5. springboot+Thymeleaf+layui 实现分页

    layui分页插件 引入相关的js和css layui:css <link rel="stylesheet" th:href="@{layui/css/layui. ...

  6. 诸神之眼 - Nmap 教程

    *注: www.heihei.work 为本人测试网站,可由 官方测试网站 scanme.nmap.org代替. NO.11.扫描端口开放的服务nmap -T4 -A -v www.heiehi.wo ...

  7. ft2000安装tigervnc

    apt update apt install tigervnc*vncserver :88 history >>history

  8. Otter远程调试

    Otter远程调试 环境配置: 机器 172.16.0.2 172.16.0.3 172.16.0.4 ZK FOLLOWER FOLLOWER LEADER MySQL Manager MySQL ...

  9. mysql3_pymysql

    python数据库编程 1.pyshon数据库接口(python DB-API) 1.为开发人员提供的数据库应用编程接口 2.python支持的数据库服务软件 mysql,oracle,sql_ser ...

  10. lvscan 查看系统中存在的所有LVM逻辑卷

    相关命令:lvresize,lvreduce,lvextend,lvdisplay,lvcreate,lvremove lvscan指令:扫描逻辑卷[语    法]lvscan [选项][功能介绍]l ...