用Auto-TensorCore代码生成优化matmul

用Auto-TensorCore代码生成优化matmul

将演示如何使用TVM Auto TensorCore CodeGen在Volta/Turing GPU上编写高性能matmul调度。这是一个透明的解决方案，可以生成大多数在ir过程中完成的转换的tensorcore内核。用户还可以编写带有tensorize的调度来生成TensorCore代码。两个解决方案使用相同的tensorcore内部函数。有关详细信息，请参阅如何使用TensorCores优化卷积资料。              准备和算法

支持2种输入数据类型：float16和int8。对于float16，累加器为float32。对于int8，累加器为int32。对于数据布局，“N”表示无转置，“T”表示转置。

import logging

import sys

import numpy as np

import tvm

from tvm import te

from tvm import autotvm

from tvm.contrib import nvcc

def matmul_nn(A, B, L, dtype="float16", layout="NN"):

k = te.reduce_axis((0, L), name="k")

if dtype == "float16":

out_type = "float"

elif dtype == "int8":

out_type = "int"

elif dtype == "int4" or dtype == "int1":

out_type = "int"

if layout == "NN":

return te.compute(

(N, M), lambda i, j: te.sum(A[i, k].astype(out_type) * B[k, j].astype(out_type), axis=k)

)

if layout == "NT":

return te.compute(

(N, M), lambda i, j: te.sum(A[k, i].astype(out_type) * B[k, j].astype(out_type), axis=k)

)

if layout == "TN":

return te.compute(

(N, M), lambda i, j: te.sum(A[i, k].astype(out_type) * B[j, k].astype(out_type), axis=k)

)

if layout == "TT":

return te.compute(

(N, M), lambda i, j: te.sum(A[k, i].astype(out_type) * B[j, k].astype(out_type), axis=k)

)

Scheduling the Computation

这个调度和GPU上的非tensorcore matmul调度没有什么不同。有关优化matmul调度的基础知识，请参考How to optimize GEMM on CPU资料。当设置“tensor_core”pragma时，“rewrite for tensorcore”ir pass将自动转换tensorcore codegen的调度，否则将生成性能较低但功能相同的普通CUDA代码。

注意

TensorCore要求

请注意，在以下两种情况下，即使设置了“tensor_core”标注，TVM仍将返回到正常的CUDA codegen：（1）输入矩阵的m、n或k不是16的倍数；（2）CUDA9上的扭曲分片大小不是{16x16x16，32x8x16，8x32x16}中的一种。

在这个项目中，storage_align用于减少共享内存的库冲突。有关storage_align原语的用法，请参阅本文档。简言之，需要为一些共享内存缓冲区添加偏移量，以减少内存冲突。根据wmma doc，load_matrix_sync的步长必须是16字节的倍数，因此选择8作为float16的偏移量，16作为int8的偏移量。

使用AutoTVM搜索此计划中的最佳配置。

@autotvm.template("tutorial/auto_tensorcore/test_gemm")

def test_gemm(N, L, M, dtype, layout):

if layout == "NN":

shape_a = (N, L)

shape_b = (L, M)

elif layout == "NT":

shape_a = (L, N)

shape_b = (L, M)

elif layout == "TN":

shape_a = (N, L)

shape_b = (M, L)

elif layout == "TT":

shape_a = (L, N)

shape_b = (M, L)

else:

print("Unsupported layout:", layout)

sys.exit(1)

A = te.placeholder(shape_a, name="A", dtype=dtype)

B = te.placeholder(shape_b, name="B", dtype=dtype)

C = matmul_nn(A, B, L, dtype, layout)

s = te.create_schedule(C.op)

y, x = s[C].op.axis

k = s[C].op.reduce_axis[0]

# storage_align params

factor = 16

offset = 8

if dtype == "int8":

factor = 32

offset = 16

elif dtype == "int4":

factor = 64

offset = 32

elif dtype == "int1":

factor = 256

offset = 128

# create cache stages

AA = s.cache_read(A, "shared", [C])

if layout == "NN" or layout == "TN":

s[AA].storage_align(AA.op.axis[0], factor, offset)

AL = s.cache_read(AA, "local", [C])

BB = s.cache_read(B, "shared", [C])

if layout == "TT" or layout == "NT":

s[BB].storage_align(BB.op.axis[0], factor, offset)

BL = s.cache_read(BB, "local", [C])

CL = s.cache_write(C, "local")

# autotvm search space definition

cfg = autotvm.get_config()

cfg.define_knob("bx", [2, 4, 8])

cfg.define_knob("by", [8, 16, 32, 64])

cfg.define_knob("step_k", [1, 2, 4, 8, 16, 32])

cfg.define_knob("v", [4, 8, 16, 32])

by = cfg["by"].val

bx = cfg["bx"].val

step_k = cfg["step_k"].val

v = cfg["v"].val

# thread tile

TX = 8

TY = 1

if dtype == "int4" or dtype == "int1":

TX = 2

# warp tile

warp_tile_m = 16  # it could also be 8 or 32 on CUDA version >= 10.0

warp_tile_k = 16  # it must be 16 for fp16/int8 data type

if dtype == "int4":

warp_tile_m = 8

warp_tile_k = 32

elif dtype == "int1":

warp_tile_m = 8

warp_tile_k = 128

# block tile

tile_x = bx * TX

tile_y = by * TY

yo, ty = s[C].split(y, tile_y)

ty, yi = s[C].split(ty, TY)

# schedule for C stage

xo, xi = s[C].split(x, tile_x)

WX = min(warp_tile_m, tile_x)

tz, xi = s[C].split(xi, WX)

tx, xi = s[C].split(xi, TX)

s[C].reorder(yo, xo, tz, ty, tx, yi, xi)

s[C].bind(yo, te.thread_axis("blockIdx.y"))

s[C].bind(xo, te.thread_axis("blockIdx.x"))

s[C].bind(ty, te.thread_axis("threadIdx.y"))

s[C].bind(tz, te.thread_axis("threadIdx.z"))

s[C].bind(tx, te.thread_axis("threadIdx.x"))

# schedule for CL stage

ko, ki = s[CL].split(k, step_k * warp_tile_k)

kl, ki = s[CL].split(ki, warp_tile_k)

s[CL].compute_at(s[C], tx)

yo, xo = CL.op.axis

s[CL].reorder(ko, kl, ki, yo, xo)

# schedule for AA stage

s[AA].compute_at(s[CL], ko)

xo, xi = s[AA].split(s[AA].op.axis[1], factor=bx * v)

tz, tx = s[AA].split(xi, factor=(WX // TX) * v)

tx, vec = s[AA].split(tx, factor=v)

fused = s[AA].fuse(s[AA].op.axis[0], xo)

_, ty = s[AA].split(fused, factor=by)

s[AA].bind(ty, te.thread_axis("threadIdx.y"))

s[AA].bind(tz, te.thread_axis("threadIdx.z"))

s[AA].bind(tx, te.thread_axis("threadIdx.x"))

# vectorization is very important for float16/int8 inputs

s[AA].vectorize(vec)

# schedule for BB stage

s[BB].compute_at(s[CL], ko)

xo, xi = s[BB].split(s[BB].op.axis[1], factor=bx * v)

tz, tx = s[BB].split(xi, factor=(WX // TX) * v)

tx, vec = s[BB].split(tx, factor=v)

fused = s[BB].fuse(s[BB].op.axis[0], xo)

_, ty = s[BB].split(fused, factor=by)

s[BB].bind(ty, te.thread_axis("threadIdx.y"))

s[BB].bind(tz, te.thread_axis("threadIdx.z"))

s[BB].bind(tx, te.thread_axis("threadIdx.x"))

s[BB].vectorize(vec)

s[AL].compute_at(s[CL], kl)

s[BL].compute_at(s[CL], kl)

# set the 'tensor_core' pragma for tensorcore codegen

s[CL].pragma(ko, "tensor_core")

return s, [A, B, C]

AutoTune and Test

最后，使用一个调谐器来调整调度，生成具有最佳配置的代码，并运行内核与numpy进行比较，以检查结果是否正确。

# check whether the gpu has tensorcore

if not tvm.gpu(0).exist or not tvm.runtime.enabled("cuda"):

raise Exception("skip building this tutorial because cuda is not enabled..")

ctx = tvm.gpu()

if not nvcc.have_tensorcore(ctx.compute_version):

raise Exception("the gpu has no tensorcore, skipping...")

M, N, L = 512, 32, 512

dtype = "float16"

layout = "NN"

if len(sys.argv) >= 4:

M, N, L = int(sys.argv[1]), int(sys.argv[2]), int(sys.argv[3])

if len(sys.argv) >= 5:

dtype = sys.argv[4]

if len(sys.argv) >= 6:

layout = sys.argv[5]

# check whether current gpu arch support support current dtype's wmma codegen

cuda_compute_capability = tvm.runtime._ffi_api.GetDeviceAttr(2, 0, 4)

major, minor = nvcc.parse_compute_version(cuda_compute_capability)

if dtype == "int8":

assert major == 7 and minor >= 2

elif dtype == "int4" or dtype == "int1":

# int4/int1 only support layout TN

assert major == 7 and minor == 5 and layout == "TN"

def tune_and_evaluate(M, N, L, dtype, layout):

task = autotvm.task.create(

"tutorial/auto_tensorcore/test_gemm", args=(N, L, M, dtype, layout), target="cuda"

)

print(task.config_space)

logging.getLogger("autotvm").setLevel(logging.DEBUG)

logging.getLogger("autotvm").addHandler(logging.StreamHandler(sys.stdout))

measure_option = autotvm.measure_option(builder="local", runner=autotvm.LocalRunner(number=5))

tuner = autotvm.tuner.XGBTuner(task)

tuner.tune(

n_trial=1000,

measure_option=measure_option,

callbacks=[autotvm.callback.log_to_file("matmul.log")],

)

dispatch_context = autotvm.apply_history_best("matmul.log")

best_config = dispatch_context.query(task.target, task.workload)

print("\nBest config:")

print(best_config)

with autotvm.apply_history_best("matmul.log"):

with tvm.target.Target("cuda"):

s, arg_bufs = test_gemm(N, L, M, dtype, layout)

print(tvm.lower(s, arg_bufs, simple_mode=True))

func = tvm.build(s, arg_bufs)

dev_module = func.imported_modules[0]

print(dev_module.get_source())

# check correctness

if layout == "NN":

shape_a = (N, L)

shape_b = (L, M)

elif layout == "NT":

shape_a = (L, N)

shape_b = (L, M)

elif layout == "TN":

shape_a = (N, L)

shape_b = (M, L)

elif layout == "TT":

shape_a = (L, N)

shape_b = (M, L)

a_np = None

b_np = None

c_np = None

c_np_type = None

if dtype == "float16":

c_np_type = np.float32

a_np = np.random.uniform(size=shape_a).astype(np.float16)

b_np = np.random.uniform(size=shape_b).astype(np.float16)

if layout == "NN":

c_np = np.dot(a_np, b_np)

elif layout == "NT":

c_np = np.dot(a_np.T, b_np)

elif layout == "TN":

c_np = np.dot(a_np, b_np.T)

elif layout == "TT":

c_np = np.dot(a_np.T, b_np.T)

elif dtype == "int8":

c_np_type = np.int32

a_np = np.random.randint(low=-128, high=127, size=shape_a).astype(np.int8)

b_np = np.random.randint(low=-128, high=127, size=shape_b).astype(np.int8)

if layout == "NN":

c_np = np.dot(a_np.astype(np.int32), b_np.astype(np.int32))

elif layout == "NT":

c_np = np.dot(a_np.astype(np.int32).T, b_np.astype(np.int32))

elif layout == "TN":

c_np = np.dot(a_np.astype(np.int32), b_np.astype(np.int32).T)

elif layout == "TT":

c_np = np.dot(a_np.astype(np.int32).T, b_np.astype(np.int32).T)

elif dtype == "int4":

c_np_type = np.int32

a_np_int = np.random.randint(low=-8, high=7, size=shape_a).astype(np.int32)

b_np_int = np.random.randint(low=-8, high=7, size=shape_b).astype(np.int32)

# "TN"

c_np = np.dot(a_np_int.astype(np.int32), b_np_int.astype(np.int32).T)

a_np = np.zeros(shape=(N, int(L / 8)), dtype=np.int32)

b_np = np.zeros(shape=(M, int(L / 8)), dtype=np.int32)

# a_np --> col_major

for i in range(N):

for j in range(int(L / 8)):

for k in range(8):

a_np[i, j] = a_np[i, j] | ((a_np_int[i, j * 8 + k] & 0xF) << ((7 - k) * 4))

# b_np --> row_major

for i in range(M):

for j in range(int(L / 8)):

for k in range(8):

b_np[i, j] = b_np[i, j] | ((b_np_int[i, j * 8 + k] & 0xF) << ((7 - k) * 4))

elif dtype == "int1":

c_np_type = np.int32

a_np_int = np.random.randint(low=0, high=1, size=shape_a).astype(np.int32)

b_np_int = np.random.randint(low=0, high=1, size=shape_b).astype(np.int32)

# "TN"

c_np = np.dot(a_np_int.astype(np.int32), b_np_int.astype(np.int32).T)

a_np = np.zeros(shape=(N, int(L / 32)), dtype=np.int32)

b_np = np.zeros(shape=(M, int(L / 32)), dtype=np.int32)

for i in range(N):

for j in range(int(L / 32)):

for k in range(32):

a_np[i, j] = a_np[i, j] | ((a_np_int[i, j * 32 + k] & 0xF) << (31 - k))

for i in range(M):

for j in range(int(L / 32)):

for k in range(32):

b_np[i, j] = b_np[i, j] | ((b_np_int[i, j * 32 + k] & 0xF) << (31 - k))

c_tvm = tvm.nd.array(np.zeros(c_np.shape, dtype=c_np_type), ctx=ctx)

a_tvm = tvm.nd.array(a_np, ctx=ctx)

b_tvm = tvm.nd.array(b_np, ctx=ctx)

func(a_tvm, b_tvm, c_tvm)

tvm.testing.assert_allclose(c_np, c_tvm.asnumpy(), rtol=1e-3)

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

print("Time cost of this operator: %f" % evaluator(a_tvm, b_tvm, c_tvm).mean)

# We do not run the tuning in our webpage server since it takes some time.

# Uncomment the following line to run it by yourself.

# tune_and_evaluate(M, N, L, dtype, layout)

Sample Output

Best config:

[('bx', 4), ('by', 32), ('step_k', 16), ('v', 8)],,None,40

Finish loading 162 records

produce compute {

// attr [iter_var(blockIdx.y, , blockIdx.y)] thread_extent = 1

// attr [compute.local] storage_scope = "wmma.accumulator"

allocate compute.local[float32 * 256]

// attr [A.shared] storage_scope = "shared"

allocate A.shared[float16 * 8448]

// attr [B.shared] storage_scope = "shared"

allocate B.shared[float16 * 8192]

// attr [A.shared.local] storage_scope = "wmma.matrix_b"

allocate A.shared.local[float16 * 256]

// attr [B.shared.local] storage_scope = "wmma.matrix_a"

allocate B.shared.local[float16 * 256]

// attr [iter_var(blockIdx.x, , blockIdx.x)] thread_extent = 16

// attr [iter_var(threadIdx.z, , threadIdx.z)] thread_extent = 2

// attr [iter_var(threadIdx.y, , threadIdx.y)] thread_extent = 32

// attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 2

produce compute.local {

for (j.c.init, 0, 1) {

tvm_fill_fragment(compute.local, 16, 16, 16, 0, 0f)

}

// attr [iter_var(k.outer, )] pragma_tensor_core = 1

for (k.outer, 0, 2) {

produce A.shared {

for (ax0.ax1.outer.fused.outer, 0, 8) {

// attr [iter_var(threadIdx.y, , threadIdx.y)] thread_extent = 32

// attr [iter_var(threadIdx.z, , threadIdx.z)] thread_extent = 2

// attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 2

A.shared[ramp((((((ax0.ax1.outer.fused.outer*1056) + (floordiv(threadIdx.y, 8)*264)) + (floormod(threadIdx.y, 8)*32)) + (threadIdx.z*16)) + (threadIdx.x*8)), 1, 8)] = A[ramp(((((((ax0.ax1.outer.fused.outer*2048) + (floordiv(threadIdx.y, 8)*512)) + (k.outer*256)) + (floormod(threadIdx.y, 8)*32)) + (threadIdx.z*16)) + (threadIdx.x*8)), 1, 8)]

}

}

produce B.shared {

for (ax0.ax1.outer.fused.outer, 0, 8) {

// attr [iter_var(threadIdx.y, , threadIdx.y)] thread_extent = 32

// attr [iter_var(threadIdx.z, , threadIdx.z)] thread_extent = 2

// attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 2

B.shared[ramp(((((ax0.ax1.outer.fused.outer*1024) + (threadIdx.y*32)) + (threadIdx.z*16)) + (threadIdx.x*8)), 1, 8)] = B[ramp(((((((k.outer*131072) + (ax0.ax1.outer.fused.outer*16384)) + (threadIdx.y*512)) + (blockIdx.x*32)) + (threadIdx.z*16)) + (threadIdx.x*8)), 1, 8)]

}

}

for (k.inner.outer, 0, 16) {

produce A.shared.local {

for (ax1, 0, 1) {

tvm_load_matrix_sync(A.shared.local, 16, 16, 16, 0, &(A.shared[(((threadIdx.y/16)*4224) + (k.inner.outer*16))]), 264, "col_major")

}

}

produce B.shared.local {

for (ax0, 0, 1) {

for (ax1, 0, 1) {

tvm_load_matrix_sync(B.shared.local, 16, 16, 16, 0, &(B.shared[((k.inner.outer*512) + (threadIdx.z*16))]), 32, "col_major")

}

}

}

for (k.inner.inner, 0, 1) {

for (j.c, 0, 1) {

tvm_mma_sync(compute.local, 0, B.shared.local, 0, A.shared.local, 0, compute.local, 0)

}

}

}

}

}

for (j.inner.inner.inner, 0, 1) {

tvm_store_matrix_sync(compute.local, 16, 16, 16, 0, &(compute[((((threadIdx.y/16)*8192) + (blockIdx.x*32)) + (threadIdx.z*16))]), 512, "col_major")

}

}

#include <cuda_fp16.h>

__device__ half max(const half a, const half b)

{

return __hgt(__half(a), __half(b)) ? a : b;

}

__device__ half min(const half a, const half b)

{

return __hlt(__half(a), __half(b)) ? a : b;

}

__device__ half operator+(const volatile __half &a,  const volatile __half &b)

{

return __hadd(a, b);

}

__device__ half operator<=(const volatile __half &a,  const volatile __half &b)

{

return __hlt(a, b);

}

__device__ half operator*(const volatile __half &a,  const volatile __half &b)

{

return __hmul(a, b);

}

#include <mma.h>

extern "C" __global__ void default_function_kernel0( half* __restrict__ A,  half* __restrict__ B,  float* __restrict__ compute) {

nvcuda::wmma::fragment<nvcuda::wmma::accumulator, 16, 16, 16, float> compute_local[1];

__shared__ half A_shared[8448];

__shared__ half B_shared[8192];

nvcuda::wmma::fragment<nvcuda::wmma::matrix_b, 16, 16, 16, half, nvcuda::wmma::col_major> A_shared_local[1];

nvcuda::wmma::fragment<nvcuda::wmma::matrix_a, 16, 16, 16, half, nvcuda::wmma::col_major> B_shared_local[1];

for (int j_c_init = 0; j_c_init < 1; ++j_c_init) {

(void)nvcuda::wmma::fill_fragment(compute_local[0], 0.000000e+00f);

}

for (int k_outer = 0; k_outer < 2; ++k_outer) {

__syncthreads();

for (int ax0_ax1_outer_fused_outer = 0; ax0_ax1_outer_fused_outer < 8; ++ax0_ax1_outer_fused_outer) {

((__shared__ float4*)(A_shared + (((((ax0_ax1_outer_fused_outer * 1056) + ((((int)threadIdx.y) >> 3) * 264)) + ((((int)threadIdx.y) & 7) * 32)) + (((int)threadIdx.z) * 16)) + (((int)threadIdx.x) * 8))))[0] = (( float4*)(A + ((((((ax0_ax1_outer_fused_outer * 2048) + ((((int)threadIdx.y) >> 3) * 512)) + (k_outer * 256)) + ((((int)threadIdx.y) & 7) * 32)) + (((int)threadIdx.z) * 16)) + (((int)threadIdx.x) * 8))))[0];

}

for (int ax0_ax1_outer_fused_outer1 = 0; ax0_ax1_outer_fused_outer1 < 8; ++ax0_ax1_outer_fused_outer1) {

((__shared__ float4*)(B_shared + ((((ax0_ax1_outer_fused_outer1 * 1024) + (((int)threadIdx.y) * 32)) + (((int)threadIdx.z) * 16)) + (((int)threadIdx.x) * 8))))[0] = (( float4*)(B + ((((((k_outer * 131072) + (ax0_ax1_outer_fused_outer1 * 16384)) + (((int)threadIdx.y) * 512)) + (((int)blockIdx.x) * 32)) + (((int)threadIdx.z) * 16)) + (((int)threadIdx.x) * 8))))[0];

}

__syncthreads();

for (int k_inner_outer = 0; k_inner_outer < 16; ++k_inner_outer) {

for (int ax1 = 0; ax1 < 1; ++ax1) {

(void)nvcuda::wmma::load_matrix_sync(A_shared_local[0], &(A_shared[(((((int)threadIdx.y) / 16) * 4224) + (k_inner_outer * 16))]), 264);

}

for (int ax0 = 0; ax0 < 1; ++ax0) {

for (int ax11 = 0; ax11 < 1; ++ax11) {

(void)nvcuda::wmma::load_matrix_sync(B_shared_local[0], &(B_shared[((k_inner_outer * 512) + (((int)threadIdx.z) * 16))]), 32);

}

}

for (int k_inner_inner = 0; k_inner_inner < 1; ++k_inner_inner) {

for (int j_c = 0; j_c < 1; ++j_c) {

(void)nvcuda::wmma::mma_sync(compute_local[0], B_shared_local[0], A_shared_local[0], compute_local[0]);

}

}

}

}

for (int j_inner_inner_inner = 0; j_inner_inner_inner < 1; ++j_inner_inner_inner) {

(void)nvcuda::wmma::store_matrix_sync(&(compute[((((((int)threadIdx.y) / 16) * 8192) + (((int)blockIdx.x) * 32)) + (((int)threadIdx.z) * 16))]), compute_local[0], 512, nvcuda::wmma::mem_col_major);

}

}

Time cost of this operator: 0.000008

Summary

演示如何使用TVM的AutoSensorCoreCodegen生成tensorcore内核。

下载Python源代码：opt_matmul_auto_tensorcore.py

下载Jupyter笔记：opt\u matmul_auto_tensorcore.ipynb网站

https://tvm.apache.org/docs/tutorials/optimize/opt_matmul_auto_tensorcore.html#sample-output

Download Python source code: opt_matmul_auto_tensorcore.py

Download Jupyter notebook: opt_matmul_auto_tensorcore.ipynb