如何在CPU上优化GEMM（下）

Array Packing

另一个重要的技巧是数组打包。这个技巧是对数组的存储维度进行重新排序，将某个维度上的连续访问模式在平滑后转换为顺序模式。

如上图所示，在阻塞计算之后，可以观察到B的数组访问模式（扁平化后），它是规则的但不连续的。期望经过一些转换，可以得到连续访问模式。可以将[16][16]数组重新排序为[16/4][16][4]数组，这样当从压缩数组中获取相应的值时，B的访问模式将是顺序的。

We have to re-write the algorithm slightly.

packedB = te.compute((N / bn, K, bn), lambda x, y, z: B[y, x * bn + z], name=“packedB”)

C = te.compute((M, N),
lambda x, y: te.sum(A[x, k] * packedB[y // bn, k, tvm.tir.indexmod(y, bn)], axis=k),

name=“C”,

)

s = te.create_schedule(C.op)

xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)

(k,) = s[C].op.reduce_axis

ko, ki = s[C].split(k, factor=4)

s[C].reorder(xo, yo, ko, xi, ki, yi)

s[C].vectorize(yi)

x, y, z = s[packedB].op.axis

s[packedB].vectorize(z)

s[packedB].parallel(x)

func = tvm.build(s, [A, B, C], target=target, name=“mmult”)

assert func

c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), ctx)

func(a, b, c)

tvm.testing.assert_allclose(c.asnumpy(), answer,
rtol=1e-5)

evaluator = func.time_evaluator(func.entry_name, ctx,
number=10)

print(“Opt4: %f” % evaluator(a, b, c).mean)

Out:

Opt4: 0.105409

Here is the generated IR after array packing.

print(tvm.lower(s, [A, B, C], simple_mode=True))

Out:

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, 1024], []),

         B: Buffer(B_2: Pointer(float32),

float32, [1024, 1024], []),

         A: Buffer(A_2: Pointer(float32),

float32, [1024, 1024], [])}

buffer_map = {A_1: A, B_1: B, C_1: C} {

attr [packedB: Pointer(float32)] “storage_scope” =
“global”;

allocate(packedB, float32x32, [32768]) {

for (x: int32, 0, 32) “parallel” {

for (y: int32, 0, 1024) {

packedB[ramp(((x32768) + (y32)), 1,32)] = (float32x32*)B_2[ramp(((y1024) + (x32)), 1, 32)]

}

for (x.outer: int32, 0, 32) {

for (y.outer: int32, 0, 32) {

    for (x.inner.init: int32, 0, 32) {C_2[ramp((((x.outer*32768) +

(x.inner.init1024)) + (y.outer32)), 1, 32)] = broadcast(0f32, 32)

    }for (k.outer: int32, 0, 256) {for (x.inner: int32, 0, 32) {for (k.inner: int32, 0, 4) {C_2[ramp((((x.outer*32768) +

(x.inner1024)) + (y.outer32)), 1, 32)] =
((float32x32*)C_2[ramp((((x.outer32768) + (x.inner1024)) + (y.outer32)), 1,
32)] + (broadcast((float32)A_2[((((x.outer32768) + (x.inner1024)) +
(k.outer4)) + k.inner)], 32)(float32x32*)packedB[ramp((((y.outer32768) +
(k.outer128)) + (k.inner*32)), 1, 32)]))

}}}

}

Write cache for blocks

分块后，程序将结果逐块写入C，访问模式不是顺序的。因此，可以使用一个顺序缓存数组来保存块结果，并在所有块结果就绪时写入C。

s = te.create_schedule(C.op)

Allocate write cache

CC = s.cache_write(C, “global”)

xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)

Write cache is computed at yo

s[CC].compute_at(s[C], yo)

New inner axes

xc, yc = s[CC].op.axis

(k,) = s[CC].op.reduce_axis

ko, ki = s[CC].split(k, factor=4)

s[CC].reorder(ko, xc, ki, yc)

s[CC].unroll(ki)

s[CC].vectorize(yc)

x, y, z = s[packedB].op.axis

s[packedB].vectorize(z)

s[packedB].parallel(x)

func = tvm.build(s, [A, B, C], target=target, name=“mmult”)

assert func

c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), ctx)

func(a, b, c)

tvm.testing.assert_allclose(c.asnumpy(), answer,
rtol=1e-5)

evaluator = func.time_evaluator(func.entry_name, ctx, number=10)

print(“Opt5: %f” % evaluator(a, b, c).mean)

Out:

Opt5: 0.098048

Here is the generated IR after blocking.

print(tvm.lower(s, [A, B, C], simple_mode=True))

Out:

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, 1024], []),

         B: Buffer(B_2: Pointer(float32),

float32, [1024, 1024], []),

         A: Buffer(A_2: Pointer(float32),

float32, [1024, 1024], [])}

buffer_map = {A_1: A, B_1: B, C_1: C} {

attr [packedB: Pointer(float32)] “storage_scope” =
“global”;

allocate(packedB, float32x32, [32768]);

attr [C.global: Pointer(float32)] “storage_scope” =
“global”;

allocate(C.global, float32, [1024]) {

for (x: int32, 0, 32) “parallel” {

for (y: int32, 0, 1024) {

    packedB[ramp(((x*32768) + (y*32)), 1, 32)] = (float32x32*)B_2[ramp(((y*1024) + (x*32)), 1, 32)]

}

for (x.outer: int32, 0, 32) {

for (y.outer: int32, 0, 32) {

    for (x.c.init: int32, 0, 32) {C.global[ramp((x.c.init*32), 1, 32)]

= broadcast(0f32, 32)

for (k.outer: int32, 0, 256) {

      for (x.c: int32, 0, 32) {C.global[ramp((x.c*32), 1, 32)] =

((float32x32*)C.global[ramp((x.c32), 1, 32)] +
(broadcast((float32)A_2[(((x.outer32768) + (x.c1024)) + (k.outer4))],
32)(float32x32*)packedB[ramp(((y.outer32768) + (k.outer128)), 1, 32)]))

        C.global[ramp((x.c*32), 1, 32)] =

((float32x32*)C.global[ramp((x.c32), 1, 32)] +
(broadcast((float32)A_2[((((x.outer32768) + (x.c1024)) + (k.outer4)) + 1)],
32)(float32x32*)packedB[ramp((((y.outer32768) + (k.outer128)) + 32), 1,
32)]))

        C.global[ramp((x.c*32), 1, 32)] =

((float32x32*)C.global[ramp((x.c32), 1, 32)] +
(broadcast((float32)A_2[((((x.outer32768) + (x.c1024)) + (k.outer4)) + 2)],
32)(float32x32*)packedB[ramp((((y.outer32768) + (k.outer128)) + 64), 1,
32)]))

        C.global[ramp((x.c*32), 1, 32)] =

((float32x32*)C.global[ramp((x.c32), 1, 32)] +
(broadcast((float32)A_2[((((x.outer32768) + (x.c1024)) + (k.outer4)) + 3)],
32)(float32x32*)packedB[ramp((((y.outer32768) + (k.outer128)) + 96), 1,
32)]))

      }}for (x.inner: int32, 0, 32) {for (y.inner: int32, 0, 32) {C_2[((((x.outer*32768) +

(x.inner1024)) + (y.outer32)) + y.inner)] = (float32*)C.global[((x.inner*32)

y.inner)]
```
    }}
```

}

Parallel

此外，还可以利用多核处理器来实现线程级的并行化。

s = te.create_schedule(C.op)

CC = s.cache_write(C, “global”)

xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)

s[CC].compute_at(s[C], yo)

xc, yc = s[CC].op.axis

(k,) = s[CC].op.reduce_axis

ko, ki = s[CC].split(k, factor=4)

s[CC].reorder(ko, xc, ki,
yc)

s[CC].unroll(ki)

s[CC].vectorize(yc)

parallel

s[C].parallel(xo)

x, y, z = s[packedB].op.axis

s[packedB].vectorize(z)

s[packedB].parallel(x)

func = tvm.build(s, [A, B, C], target=target, name=“mmult”)

assert func

c = tvm.nd.array(numpy.zeros((M, N), dtype=dtype), ctx)

func(a, b, c)

tvm.testing.assert_allclose(c.asnumpy(), answer,
rtol=1e-5)

evaluator = func.time_evaluator(func.entry_name, ctx,
number=50)

opt6_time = evaluator(a, b, c).mean

print(“Opt6: %f” % opt6_time)

Out:

Opt6: 0.032347

Here is the generated IR after parallelization.

print(tvm.lower(s, [A, B, C], simple_mode=True))

Out:

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, 1024], []),

         B: Buffer(B_2: Pointer(float32),

float32, [1024, 1024], []),

         A: Buffer(A_2: Pointer(float32),

float32, [1024, 1024], [])}

buffer_map = {A_1: A, B_1: B, C_1: C} {

attr [packedB: Pointer(float32)] “storage_scope” =
“global”;

allocate(packedB, float32x32, [32768]) {

for (x: int32, 0, 32) “parallel” {

for (y: int32, 0, 1024) {

    packedB[ramp(((x*32768) + (y*32)), 1, 32)] = (float32x32*)B_2[ramp(((y*1024) + (x*32)), 1, 32)]

}

for (x.outer: int32, 0, 32) “parallel” {

attr [C.global: Pointer(float32)] “storage_scope” =
“global”;

allocate(C.global, float32, [1024]);

for (y.outer: int32, 0, 32) {

    for (x.c.init: int32, 0, 32) {C.global[ramp((x.c.init*32), 1, 32)]

= broadcast(0f32, 32)

    }for (k.outer: int32, 0, 256) {for (x.c: int32, 0, 32) {C.global[ramp((x.c*32), 1, 32)] =

((float32x32*)C.global[ramp((x.c32), 1, 32)] +
(broadcast((float32)A_2[(((x.outer32768) + (x.c1024)) + (k.outer4))],
32)(float32x32*)packedB[ramp(((y.outer32768) + (k.outer128)), 1, 32)]))

        C.global[ramp((x.c*32), 1, 32)] =

        C.global[ramp((x.c*32), 1, 32)] =

        C.global[ramp((x.c*32), 1, 32)] =

      }}for (x.inner: int32, 0, 32) {for (y.inner: int32, 0, 32) {C_2[((((x.outer*32768) +

(x.inner1024)) + (y.outer32)) + y.inner)] = (float32*)C.global[((x.inner*32)

y.inner)]
```
    }}
```

}

Summary

在用18行代码应用上述简单的优化之后，生成的代码可以达到MKL的60%的numpy性能。请注意，网页上的输出反映了非独占Docker容器上的运行时间，因此是不可靠的。强烈建议自己来完成，以观察TVM所获得的性能提升。

https://tvm.apache.org/docs/tutorials/optimize/opt_gemm.html#sphx-glr-tutorials-optimize-opt-gemm-py

如何在CPU上优化GEMM（下）相关推荐

如何在CPU上优化GEMM矩阵乘法
如何在CPU上优化GEMM矩阵乘法 How to optimize GEMM on CPU (TL;DR) TVM 提供抽象接口,允许用户分别描述算法和算法的实现组织(所谓的调度).通常,在高性能调度 ...
如何在 CPU 上优化 GEMM
如何在 CPU 上优化 GEMM (TL;DR) TVM 提供抽象接口,允许用户分别描述算法和算法的实施组织(所谓的调度).通常,在高性能调度中编写算法,会破坏算法的可读性和模块化.尝试各种看似有前途 ...
如何在CPU上优化GEMM（上）
如何在CPU上优化GEMM(上) (TL:DR)TVM提供了抽象接口,用户分别描述算法和算法的实现组织(所谓的调度).通常,在高性能调度中编写算法会破坏算法的可读性和模块性.尝试各种看似有希望的时间表 ...
如何在 GPU 上优化卷积
如何在 GPU 上优化卷积将演示如何在 TVM 中编写高性能卷积实现.正方形大小的输入张量和过滤器为例,假设卷积的输入具有大batch批量.在这个例子中,使用不同的布局存储数据,实现更好的数据局部性 ...
如何在GPU上优化卷积
如何在GPU上优化卷积本文将演示如何在TVM中编写高性能的卷积实现.以平方大小的输入张量和滤波器为例,并假设卷积的输入量很大.使用不同的布局来存储数据,以实现更好的数据局部性.缓冲区布局为HWCN, ...
赛灵思运行linux,玩转赛灵思Zedboard开发板（6）:如何在Zedboard上运行linux下的应用程序？...
描述电子发烧友网讯:ZedBoard开发板上的Zynq是一个ARM PS(processing system, 双核A9 + 存储管理 + 外设)+ PL(programable Logic) 结构 ...
java android 堆栈_如何在Android上的JNI下捕获SIGSEGV（分段错误）并获得堆栈跟踪？...
ITMISS 从Jelly Bean开始,您无法获取堆栈跟踪,因为READ_LOGS走了.:-(实际上,我的信号处理程序在工作时没有做任何多余的事情,并且已经发布了使用它的代码,您可以在github上 ...
linux多串口改波特率,如何在S3C2440上linux操作系统下将串口的波特率提高以致921600...
首先声明,CSDN中藏龙卧虎,说不定就出来个大牛,所以请各位大牛批评指正在说说我做的这个事情,其实听起来很简单,就是把串口的波特率提上去,硬件环境呢,就首先声明,CSDN中藏龙卧虎,说不定就出来个 ...
【YOLOv5】LabVIEW+OpenVINO让你的YOLOv5在CPU上飞起来
文章目录前言一.OpenVINO是什么二.LabVIEW视觉工具包下载与配置 1.视觉工具包的下载安装 2.OpenVINO toolkit下载安装三.模型获取四.LabVIEW+OpenV ...

如何在CPU上优化GEMM（下）

We have to re-write the algorithm slightly.

Allocate write cache

Write cache is computed at yo

New inner axes

parallel

如何在CPU上优化GEMM（下）相关推荐

最新文章

热门文章