使用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. 怎样用SQL修改某个字段的部分内容

    方法:update dede_addonarticle set body = replace(body,'#p#分页标题#e#',' ') where body like '%#p#分页标题#e#%'

  2. UVA11427玩纸牌(全概率+递推)

    题意:       一个人玩纸牌游戏,他每天最多玩n局,枚举获胜的概率是a/b,每天玩牌只要获胜概率达到p,那么他今天就不玩了,明天接着玩,如果有一天他的概率没有达到p,(没有达到p的话他今天一定是玩 ...

  3. Win64 驱动内核编程-5.内核里操作文件

    内核里操作文件 RING0 操作文件和 RING3 操作文件在流程上没什么大的区别,也是"获得文件句柄->读/写/删/改->关闭文件句柄"的模式.当然了,只能用内核 A ...

  4. Docker Swarm删除节点

    节点上的主机如果想离开的话,可以自己直接执行docker swarm leave 然后你可以发现,原本跑在自己上面的容器被转移到别的容器上了.此时如果在manager节点上docker node ls ...

  5. 【JavaScript】【KMP】Leetcode每日一题-实现strStr()

    [JavaScript]Leetcode每日一题-实现strStr() [题目描述] 实现 strStr() 函数. 给你两个字符串 haystack 和 needle ,请你在 haystack 字 ...

  6. 多线程-5.JMM之happens-before原则

    a happens-before b 翻译为a操作对b操作是可见的.可见即是指共享变量的更改能获知. 特性:传递性 原则:volatile定义的变量 写操作 happens-before 读操作 同一 ...

  7. CRM客户关系管理系统有哪些优缺点?

    CRM系统不仅仅是一种技术,也是面向企业的客户管理系统.客户关系管理软件可以帮助销售员快速地找到客户信息,帮助销售员跟踪客户直到完成订单.为提高企业销售效率,CRM被越来越多的企业所采用. 那么,作为 ...

  8. spring-第三章-jdbc

    一,回顾 aop:面向切面编程,就是将一些和主业务流程没有关系的公共代码,提取封装到切面类,通过切入点规则,可以对目标方法进行功能增强;也就是可以再目标方法执行的前后添加一段额外逻辑代码; 二,Jdb ...

  9. GIT·版本回退的指令

    阅文时长 | 0.46分钟 字数统计 | 499.2字符 主要内容 | 1.引言&背景 2.指令集合&示例 3.声明与参考资料 『GIT·版本回退的指令』 编写人 | SCscHero ...

  10. cent 7 识别exfat

    cent 7 识别exfat install fuse-exfat exfat-utils