如何使用 TensorCores 优化卷积
备注
单击 此处 下载完整的示例代码
作者：Siyuan Feng
本教程演示如何在 TVM 中使用 TensorCores 编写高性能卷积调度。在这个例子中，会假设卷积输入的 batch 较大。强烈建议前置讲解 如何在 GPU 上优化卷积。
TensorCore 介绍
每个 Tensor Core 都提供了一个 4x4x4 矩阵处理数组，它使得 D = A * B + C，其中 A、B、C 和 D 是 4x4 矩阵，如图所示。矩阵乘法输入 A 和 B 是 FP16 矩阵，而累加矩阵 C 和 D 可以是 FP16 或 FP32 矩阵。
但是，CUDA 开发者只能使用 warp 级原语 wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag) 在张量核上执行 16x16x16 半精度矩阵乘法。调用矩阵乘法之前，开发者必须使用原始 wmma::load_matrix_sync 显式地将数据从内存加载到寄存器中。NVCC 编译器将该原语转换为多个内存加载指令。在运行时，每个线程从矩阵 A 加载 16 个元素，从 B 加载 16 个元素。
准备和算法
对具有 256 个通道和 14 x 14 维度的输入张量使用固定大小。batch size 为 256，卷积过滤器包含 512 个大小为 3 x 3 的过滤器，使用 stride size 为 1 和 padding size 为 1 进行卷积。在示例中，使用 NHWCnc 内存布局。以下代码定义了 TVM 中的卷积算法。
import tvm
from tvm import te
import numpy as np
from tvm.contrib import nvcc
# 输入和过滤器的大小
batch_size = 256
height = 14
width = 14
in_channels = 256
out_channels = 512
kernel_h = 3
kernel_w = 3
pad_h = 1
pad_w = 1
stride_h = 1
stride_w = 1
# TensorCore shape
block_size = 16
assert batch_size % block_size == 0
assert in_channels % block_size == 0
assert out_channels % block_size == 0
# 输入特征图：（N，H，W，IC，n，ic）
data_shape = (
batch_size // block_size,
height,
width,
in_channels // block_size,
block_size,
block_size,
)
# Kernel: (H, W, IC, OC, ic, oc)
kernel_shape = (
kernel_h,
kernel_w,
in_channels // block_size,
out_channels // block_size,
block_size,
block_size,
)
# 输出特征图：(N, H, W, OC, n, oc)
output_shape = (
batch_size // block_size,
height,
width,
out_channels // block_size,
block_size,
block_size,
)
# Reduction axes
kh = te.reduce_axis((0, kernel_h), name="kh")
kw = te.reduce_axis((0, kernel_w), name="kw")
ic = te.reduce_axis((0, in_channels // block_size), name="ic")
ii = te.reduce_axis((0, block_size), name="ii")
# 算法
A = te.placeholder(data_shape, name="A", dtype="float16")
W = te.placeholder(kernel_shape, name="W", dtype="float16")
Apad = te.compute(
(
batch_size // block_size,
height + 2 * pad_h,
width + 2 * pad_w,
in_channels // block_size,
block_size,
block_size,
),
lambda n, h, w, i, nn, ii: tvm.tir.if_then_else(
tvm.tir.all(h >= pad_h, h - pad_h < height, w >= pad_w, w - pad_w < width),
A[n, h - pad_h, w - pad_w, i, nn, ii],
tvm.tir.const(0.0, "float16"),
),
name="Apad",
)
Conv = te.compute(
output_shape,
lambda n, h, w, o, nn, oo: te.sum(
Apad[n, h * stride_h + kh, w * stride_w + kw, ic, nn, ii].astype("float32")
* W[kh, kw, ic, o, ii, oo].astype("float32"),
axis=[ic, kh, kw, ii],
),
name="Conv",
)
s = te.create_schedule(Conv.op)
s[Apad].compute_inline()
内存范围
传统的 GPU 调度有全局、共享和本地内存范围。为了支持 TensorCores，添加另外三个特殊的内存范围：wmma.matrix_a, wmma.matrix_b 和 wmma.accumulator。在硬件上，所有片段范围都存储在芯片上寄存器级别，与本地内存相同。
# 指定内存层次结构
AS = s.cache_read(Apad, "shared", [Conv])
WS = s.cache_read(W, "shared", [Conv])
AF = s.cache_read(AS, "wmma.matrix_a", [Conv])
WF = s.cache_read(WS, "wmma.matrix_b", [Conv])
ConvF = s.cache_write(Conv, "wmma.accumulator")
定义张量内联函数
实际上，TensorCore 是一种特殊的硬件操作。因此，可以只用 tensorize 将一个计算单元替换为 TensorCore 指令。首先需要定义张量内联函数。
TensorCore 中有四个基本操作：fill_fragment， load_matrix， mma_sync 和 store_matrix。由于 fill_fragment 和 mma_sync 都用于矩阵乘法，所以可以只写以下三个内联函数。
def intrin_wmma_load_matrix(scope):
n = 16
A = te.placeholder((n, n), name="A", dtype="float16")
BA = tvm.tir.decl_buffer(A.shape, A.dtype, scope="shared", data_alignment=32, offset_factor=256)
C = te.compute((n, n), lambda i, j: A[i, j], name="C")
BC = tvm.tir.decl_buffer(C.shape, C.dtype, scope=scope, data_alignment=32, offset_factor=256)
def intrin_func(ins, outs):
ib = tvm.tir.ir_builder.create()
BA = ins[0]
BC = outs[0]
ib.emit(
tvm.tir.call_intrin(
"handle",
"tir.tvm_load_matrix_sync",
BC.data,
n,
n,
n,
BC.elem_offset // 256,
BA.access_ptr("r"),
n,
"row_major",
)
)
return ib.get()
return te.decl_tensor_intrin(C.op, intrin_func, binds={A: BA, C: BC})
def intrin_wmma_gemm():
n = 16
A = te.placeholder((n, n), name="A", dtype="float16")
B = te.placeholder((n, n), name="B", dtype="float16")
k = te.reduce_axis((0, n), name="k")
C = te.compute(
(n, n),
lambda ii, jj: te.sum(A[ii, k].astype("float") * B[k, jj].astype("float"), axis=k),
name="C",
)
BA = tvm.tir.decl_buffer(
A.shape, A.dtype, name="BA", scope="wmma.matrix_a", data_alignment=32, offset_factor=256
)
BB = tvm.tir.decl_buffer(
B.shape, B.dtype, name="BB", scope="wmma.matrix_b", data_alignment=32, offset_factor=256
)
BC = tvm.tir.decl_buffer(
C.shape, C.dtype, name="BC", scope="wmma.accumulator", data_alignment=32, offset_factor=256
)
def intrin_func(ins, outs):
BA, BB = ins
(BC,) = outs
def init():
ib = tvm.tir.ir_builder.create()
ib.emit(
tvm.tir.call_intrin(
"handle", "tir.tvm_fill_fragment", BC.data, n, n, n, BC.elem_offset // 256, 0.0
)
)
return ib.get()
def update():
ib = tvm.tir.ir_builder.create()
ib.emit(
tvm.tir.call_intrin(
"handle",
"tir.tvm_mma_sync",
BC.data,
BC.elem_offset // 256,
BA.data,
BA.elem_offset // 256,
BB.data,
BB.elem_offset // 256,
BC.data,
BC.elem_offset // 256,
)
)
return ib.get()
return update(), init(), update()
return te.decl_tensor_intrin(C.op, intrin_func, binds={A: BA, B: BB, C: BC})
def intrin_wmma_store_matrix():
n = 16
A = te.placeholder((n, n), name="A", dtype="float32")
BA = tvm.tir.decl_buffer(
A.shape, A.dtype, scope="wmma.accumulator", data_alignment=32, offset_factor=256
)
C = te.compute((n, n), lambda i, j: A[i, j], name="C")
BC = tvm.tir.decl_buffer(C.shape, C.dtype, scope="global", data_alignment=32, offset_factor=256)
def intrin_func(ins, outs):
ib = tvm.tir.ir_builder.create()
BA = ins[0]
BC = outs[0]
ib.emit(
tvm.tir.call_intrin(
"handle",
"tir.tvm_store_matrix_sync",
BA.data,
n,
n,
n,
BA.elem_offset // 256,
BC.access_ptr("w"),
n,
"row_major",
)
)
return ib.get()
return te.decl_tensor_intrin(C.op, intrin_func, binds={A: BA, C: BC})
调度计算
要在 TVM 中使用 TensorCores，必须将计算调度到特定的结构中，从而匹配张量内联函数。和传统的 GPU 程序一样，也可以用共享内存来提升速度。如果对分块和共享内存有任何疑问，请参阅 如何在 GPU 上优化卷积。
在这个例子中，每个块包含 2x4 个 warp，每个 warp 调用 4x2 TensorCore 指令。因此，每个 warp 的输出 shape 为 64x32，每个块输出 128x128 个 titles。由于共享内存空间的限制，一次只加载 2 个块（2x128x128 tiles）。
备注
Warp 级操作
所有 TensorCore 指令都是 warp 级指令，这意味着一个 warp 中的所有 32 个线程应该同时执行此指令。
threadIdx.x extent=32 是解决此问题的最简单方法之一。除了那些直接或间接包含 TensorCore 内联函数的循环，可以将 threadIdx.x 绑定到任何循环。
这并不是唯一的解决方案，唯一应该做的就是确保一个 warp 中的所有线程都可以同时调用 TensorCore。
# 定义 tile 大小
block_row_warps = 4
block_col_warps = 2
warp_row_tiles = 2
warp_col_tiles = 4
warp_size = 32
chunk = 2
block_x = te.thread_axis("blockIdx.x")
block_y = te.thread_axis("blockIdx.y")
block_z = te.thread_axis("blockIdx.z")
thread_x = te.thread_axis("threadIdx.x")
thread_y = te.thread_axis("threadIdx.y")
thread_z = te.thread_axis("threadIdx.z")
nc, hc, wc, oc, nnc, ooc = Conv.op.axis
block_k = s[Conv].fuse(hc, wc)
s[Conv].bind(block_k, block_z)
nc, nci = s[Conv].split(nc, factor=warp_row_tiles)
block_i, nc = s[Conv].split(nc, factor=block_row_warps)
oc, oci = s[Conv].split(oc, factor=warp_col_tiles)
block_j, oc = s[Conv].split(oc, factor=block_col_warps)
s[Conv].reorder(block_k, block_i, block_j, nc, oc, nci, oci, nnc, ooc)
s[Conv].bind(block_i, block_x)
s[Conv].bind(block_j, block_y)
s[Conv].bind(nc, thread_y)
s[Conv].bind(oc, thread_z)
# 调度本地计算
s[ConvF].compute_at(s[Conv], oc)
n, h, w, o, nnf, oof = ConvF.op.axis
ko, ki = s[ConvF].split(ic, factor=chunk)
s[ConvF].reorder(ko, kh, ki, kw, n, o, nnf, oof, ii)
# 将中间计算移动到每个输出计算块中
s[AF].compute_at(s[ConvF], kw)
s[WF].compute_at(s[ConvF], kw)
# A 的共享内存调度
s[AS].compute_at(s[ConvF], kh)
n, h, w, i, nn, ii = AS.op.axis
tx, xo = s[AS].split(n, nparts=block_row_warps)
ty, yo = s[AS].split(xo, nparts=block_col_warps)
t = s[AS].fuse(nn, ii)
to, ti = s[AS].split(t, factor=warp_size)
s[AS].bind(tx, thread_y)
s[AS].bind(ty, thread_z)
s[AS].bind(ti, thread_x)
# W 的共享内存调度
s[WS].compute_at(s[ConvF], kh)
kh, kw, ic, o, ii, oo = WS.op.axis
tx, xo = s[WS].split(o, nparts=block_row_warps)
ty, yo = s[WS].split(xo, nparts=block_col_warps)
t = s[WS].fuse(ii, oo)
to, ti = s[WS].split(t, nparts=warp_size)
s[WS].bind(tx, thread_y)
s[WS].bind(ty, thread_z)
s[WS].bind(to, thread_x)
s[WS].vectorize(ti)
print(tvm.lower(s, [A, W, Conv], simple_mode=True))
输出结果：
@main = primfn(A_1: handle, W_1: handle, Conv_1: handle) -> ()
attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}
buffers = {A: Buffer(A_2: Pointer(float16), float16, [12845056], []),
W: Buffer(W_2: Pointer(float16), float16, [1179648], []),
Conv: Buffer(Conv_2: Pointer(float32), float32, [25690112], [])}
buffer_map = {A_1: A, W_1: W, Conv_1: Conv}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float16, [16, 14, 14, 16, 16, 16], []), W_1: W_3: Buffer(W_2, float16, [3, 3, 16, 32, 16, 16], []), Conv_1: Conv_3: Buffer(Conv_2, float32, [16, 14, 14, 32, 16, 16], [])} {
attr [IterVar(blockIdx.z: int32, (nullptr), "ThreadIndex", "blockIdx.z")] "thread_extent" = 196;
allocate(Conv.wmma.accumulator: Pointer(wmma.accumulator float32), float32, [2048]), storage_scope = wmma.accumulator;
allocate(Apad.shared: Pointer(shared float16), float16, [12288]), storage_scope = shared;
allocate(W.shared: Pointer(shared float16), float16, [12288]), storage_scope = shared;
allocate(Apad.shared.wmma.matrix_a: Pointer(wmma.matrix_a float16), float16, [512]), storage_scope = wmma.matrix_a;
allocate(W.shared.wmma.matrix_b: Pointer(wmma.matrix_b float16), float16, [1024]), storage_scope = wmma.matrix_b;
attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = 2;
attr [IterVar(blockIdx.y: int32, (nullptr), "ThreadIndex", "blockIdx.y")] "thread_extent" = 4;
attr [IterVar(threadIdx.y: int32, (nullptr), "ThreadIndex", "threadIdx.y")] "thread_extent" = 4;
attr [IterVar(threadIdx.z: int32, (nullptr), "ThreadIndex", "threadIdx.z")] "thread_extent" = 2 {
for (n.c.init: int32, 0, 2) {
for (o.c.init: int32, 0, 4) {
for (nn.c.init: int32, 0, 16) {
for (oo.c.init: int32, 0, 16) {
Conv.wmma.accumulator_1: Buffer(Conv.wmma.accumulator, float32, [2048], [], scope="wmma.accumulator")[((((n.c.init*1024) + (o.c.init*256)) + (nn.c.init*16)) + oo.c.init)] = 0f32
}
}
}
}
for (ic.outer: int32, 0, 8) {
for (kh: int32, 0, 3) {
for (ax2: int32, 0, 3) {
for (ax3: int32, 0, 2) {
for (ax4.ax5.fused.outer: int32, 0, 8) {
let cse_var_2: int32 = (ax3*256)
let cse_var_1: int32 = (ax4.ax5.fused.outer*32)
attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 32;
Apad.shared_1: Buffer(Apad.shared, float16, [12288], [], scope="shared")[((((((threadIdx.y*3072) + (threadIdx.z*1536)) + (ax2*512)) + cse_var_2) + cse_var_1) + threadIdx.x)] = @tir.if_then_else(((((1 <= (floordiv(blockIdx.z, 14) + kh)) && ((floordiv(blockIdx.z, 14) + kh) < 15)) && (1 <= (ax2 + floormod(blockIdx.z, 14)))) && ((ax2 + floormod(blockIdx.z, 14)) < 15)), A[(((((((((((blockIdx.x*6422528) + (threadIdx.y*1605632)) + (threadIdx.z*802816)) + (kh*57344)) + (blockIdx.z*4096)) + (ax2*4096)) + (ic.outer*512)) + cse_var_2) + cse_var_1) + threadIdx.x) - 61440)], 0f16, dtype=float16)
}
}
}
for (ax1: int32, 0, 3) {
for (ax2_1: int32, 0, 2) {
attr [IterVar(threadIdx.x, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 32;
W.shared_1: Buffer(W.shared, float16, [12288], [], scope="shared")[ramp((((((ax1*4096) + (ax2_1*2048)) + (threadIdx.y*512)) + (threadIdx.z*256)) + (threadIdx.x*8)), 1, 8)] = W[ramp(((((((((kh*393216) + (ax1*131072)) + (ic.outer*16384)) + (ax2_1*8192)) + (blockIdx.y*2048)) + (threadIdx.y*512)) + (threadIdx.z*256)) + (threadIdx.x*8)), 1, 8)]
}
}
for (ic.inner: int32, 0, 2) {
for (kw: int32, 0, 3) {
for (ax0: int32, 0, 2) {
for (ax4: int32, 0, 16) {
for (ax5: int32, 0, 16) {
let cse_var_3: int32 = (ax4*16)
Apad.shared.wmma.matrix_a_1: Buffer(Apad.shared.wmma.matrix_a, float16, [512], [], scope="wmma.matrix_a")[(((ax0*256) + cse_var_3) + ax5)] = Apad.shared_1[((((((threadIdx.y*3072) + (ax0*1536)) + (kw*512)) + (ic.inner*256)) + cse_var_3) + ax5)]
}
}
}
for (ax3_1: int32, 0, 4) {
for (ax4_1: int32, 0, 16) {
for (ax5_1: int32, 0, 16) {
let cse_var_5: int32 = (ax3_1*256)
let cse_var_4: int32 = (ax4_1*16)
W.shared.wmma.matrix_b_1: Buffer(W.shared.wmma.matrix_b, float16, [1024], [], scope="wmma.matrix_b")[((cse_var_5 + cse_var_4) + ax5_1)] = W.shared_1[((((((kw*4096) + (ic.inner*2048)) + (threadIdx.z*1024)) + cse_var_5) + cse_var_4) + ax5_1)]
}
}
}
for (n.c: int32, 0, 2) {
for (o.c: int32, 0, 4) {
for (nn.c: int32, 0, 16) {
for (oo.c: int32, 0, 16) {
for (ii: int32, 0, 16) {
let cse_var_8: int32 = (o.c*256)
let cse_var_7: int32 = (nn.c*16)
let cse_var_6: int32 = ((((n.c*1024) + cse_var_8) + cse_var_7) + oo.c)
Conv.wmma.accumulator_1[cse_var_6] = (Conv.wmma.accumulator_1[cse_var_6] + (cast(float32, Apad.shared.wmma.matrix_a_1[(((n.c*256) + cse_var_7) + ii)])*cast(float32, W.shared.wmma.matrix_b_1[((cse_var_8 + (ii*16)) + oo.c)])))
}
}
}
}
}
}
}
}
}
for (n.inner: int32, 0, 2) {
for (o.inner: int32, 0, 4) {
for (nn: int32, 0, 16) {
for (oo: int32, 0, 16) {
let cse_var_10: int32 = (o.inner*256)
let cse_var_9: int32 = (nn*16)
Conv[(((((((((blockIdx.x*12845056) + (threadIdx.y*3211264)) + (n.inner*1605632)) + (blockIdx.z*8192)) + (blockIdx.y*2048)) + (threadIdx.z*1024)) + cse_var_10) + cse_var_9) + oo)] = Conv.wmma.accumulator_1[((((n.inner*1024) + cse_var_10) + cse_var_9) + oo)]
}
}
}
}
}
}
将计算降级为内联函数
最后一个阶段将计算循环降级到 TensorCore 硬件内联函数，这是通过将 2D 卷积映射到张量内联函数实现的。
s[AF].tensorize(AF.op.axis[-2], intrin_wmma_load_matrix("wmma.matrix_a"))
s[WF].tensorize(WF.op.axis[-2], intrin_wmma_load_matrix("wmma.matrix_b"))
s[Conv].tensorize(nnc, intrin_wmma_store_matrix())
s[ConvF].tensorize(nnf, intrin_wmma_gemm())
print(tvm.lower(s, [A, W, Conv], simple_mode=True))
输出结果：
@main = primfn(A_1: handle, W_1: handle, Conv_1: handle) -> ()
attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}
buffers = {A: Buffer(A_2: Pointer(float16), float16, [12845056], []),
W: Buffer(W_2: Pointer(float16), float16, [1179648], []),
Conv: Buffer(Conv_2: Pointer(float32), float32, [25690112], [])}
buffer_map = {A_1: A, W_1: W, Conv_1: Conv}
preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float16, [16, 14, 14, 16, 16, 16], []), W_1: W_3: Buffer(W_2, float16, [3, 3, 16, 32, 16, 16], []), Conv_1: Conv_3: Buffer(Conv_2, float32, [16, 14, 14, 32, 16, 16], [])} {
attr [IterVar(blockIdx.z: int32, (nullptr), "ThreadIndex", "blockIdx.z")] "thread_extent" = 196;
allocate(Conv.wmma.accumulator: Pointer(wmma.accumulator float32), float32, [2048]), storage_scope = wmma.accumulator;
allocate(Apad.shared: Pointer(shared float16), float16, [12288]), storage_scope = shared;
allocate(W.shared: Pointer(shared float16), float16, [12288]), storage_scope = shared;
allocate(Apad.shared.wmma.matrix_a: Pointer(wmma.matrix_a float16), float16, [512]), storage_scope = wmma.matrix_a;
allocate(W.shared.wmma.matrix_b: Pointer(wmma.matrix_b float16), float16, [1024]), storage_scope = wmma.matrix_b;
attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = 2;
attr [IterVar(blockIdx.y: int32, (nullptr), "ThreadIndex", "blockIdx.y")] "thread_extent" = 4;
attr [IterVar(threadIdx.y: int32, (nullptr), "ThreadIndex", "threadIdx.y")] "thread_extent" = 4;
attr [IterVar(threadIdx.z: int32, (nullptr), "ThreadIndex", "threadIdx.z")] "thread_extent" = 2 {
for (n.c.init: int32, 0, 2) {
for (o.c.init: int32, 0, 4) {
@tir.tvm_fill_fragment(Conv.wmma.accumulator, 16, 16, 16, ((n.c.init*4) + o.c.init), 0f32, dtype=handle)
}
}
for (ic.outer: int32, 0, 8) {
for (kh: int32, 0, 3) {
for (ax2: int32, 0, 3) {
for (ax3: int32, 0, 2) {
for (ax4.ax5.fused.outer: int32, 0, 8) {
let cse_var_2: int32 = (ax3*256)
let cse_var_1: int32 = (ax4.ax5.fused.outer*32)
attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 32;
Apad.shared_1: Buffer(Apad.shared, float16, [12288], [], scope="shared")[((((((threadIdx.y*3072) + (threadIdx.z*1536)) + (ax2*512)) + cse_var_2) + cse_var_1) + threadIdx.x)] = @tir.if_then_else(((((1 <= (floordiv(blockIdx.z, 14) + kh)) && ((floordiv(blockIdx.z, 14) + kh) < 15)) && (1 <= (ax2 + floormod(blockIdx.z, 14)))) && ((ax2 + floormod(blockIdx.z, 14)) < 15)), A[(((((((((((blockIdx.x*6422528) + (threadIdx.y*1605632)) + (threadIdx.z*802816)) + (kh*57344)) + (blockIdx.z*4096)) + (ax2*4096)) + (ic.outer*512)) + cse_var_2) + cse_var_1) + threadIdx.x) - 61440)], 0f16, dtype=float16)
}
}
}
for (ax1: int32, 0, 3) {
for (ax2_1: int32, 0, 2) {
attr [IterVar(threadIdx.x, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 32;
W.shared_1: Buffer(W.shared, float16, [12288], [], scope="shared")[ramp((((((ax1*4096) + (ax2_1*2048)) + (threadIdx.y*512)) + (threadIdx.z*256)) + (threadIdx.x*8)), 1, 8)] = W[ramp(((((((((kh*393216) + (ax1*131072)) + (ic.outer*16384)) + (ax2_1*8192)) + (blockIdx.y*2048)) + (threadIdx.y*512)) + (threadIdx.z*256)) + (threadIdx.x*8)), 1, 8)]
}
}
for (ic.inner: int32, 0, 2) {
for (kw: int32, 0, 3) {
for (ax0: int32, 0, 2) {
@tir.tvm_load_matrix_sync(Apad.shared.wmma.matrix_a, 16, 16, 16, ax0, @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float16), Apad.shared, ((((threadIdx.y*3072) + (ax0*1536)) + (kw*512)) + (ic.inner*256)), 256, 1, dtype=handle), 16, "row_major", dtype=handle)
}
for (ax3_1: int32, 0, 4) {
@tir.tvm_load_matrix_sync(W.shared.wmma.matrix_b, 16, 16, 16, ax3_1, @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float16), W.shared, ((((kw*4096) + (ic.inner*2048)) + (threadIdx.z*1024)) + (ax3_1*256)), 256, 1, dtype=handle), 16, "row_major", dtype=handle)
}
for (n.c: int32, 0, 2) {
for (o.c: int32, 0, 4) {
let cse_var_3: int32 = ((n.c*4) + o.c)
@tir.tvm_mma_sync(Conv.wmma.accumulator, cse_var_3, Apad.shared.wmma.matrix_a, n.c, W.shared.wmma.matrix_b, o.c, Conv.wmma.accumulator, cse_var_3, dtype=handle)
}
}
}
}
}
}
for (n.inner: int32, 0, 2) {
for (o.inner: int32, 0, 4) {
@tir.tvm_store_matrix_sync(Conv.wmma.accumulator, 16, 16, 16, ((n.inner*4) + o.inner), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), Conv_2, (((((((blockIdx.x*12845056) + (threadIdx.y*3211264)) + (n.inner*1605632)) + (blockIdx.z*8192)) + (blockIdx.y*2048)) + (threadIdx.z*1024)) + (o.inner*256)), 256, 2, dtype=handle), 16, "row_major", dtype=handle)
}
}
}
}
生成 CUDA 内核
最后使用 TVM 生成和编译 CUDA 内核，并评估卷积的延迟。由于 TensorCores 仅支持 Compute Capability 7.0 或更高版本的 NVIDIA GPU，因此可能无法在构建服务器上运行。
dev = tvm.cuda(0)
if nvcc.have_tensorcore(dev.compute_version):
with tvm.transform.PassContext(config={"tir.UnrollLoop": {"auto_max_step": 16}}):
func = tvm.build(s, [A, W, Conv], "cuda")
a_np = np.random.uniform(size=data_shape).astype(A.dtype)
w_np = np.random.uniform(size=kernel_shape).astype(W.dtype)
a = tvm.nd.array(a_np, dev)
w = tvm.nd.array(w_np, dev)
c = tvm.nd.array(np.zeros(output_shape, dtype=Conv.dtype), dev)
evaluator = func.time_evaluator(func.entry_name, dev, number=10)
print("conv2d with tensor core: %f ms" % (evaluator(a, w, c).mean * 1e3))
输出结果：
conv2d with tensor core: 6.835711 ms
总结
本教程演示如何用 TVM 调度原语在特定 GPU 上调用 TensorCore。