如何使用TensorCores優化卷積
本文將演示如何在TVM中使用TensorCores編寫高性能的卷積計劃。假設卷積的輸入有大量數據。首先介紹如何在GPU上優化卷積。
TensorCore簡介
每個Tensor核心都提供一個4x4x4的矩陣處理陣列，該陣列可以運行 ，其中A，B，C和D是4x4矩陣，如圖所示。矩陣乘法輸入A和B是FP16矩陣，而累加矩陣C和D可以是FP16或FP32矩陣。D = A * B + C
但是，CUDA程序員只能使用扭曲級原語，在張量核上執行16x16x16半精度矩陣乘法。在調用矩陣乘法之前，程序員必須將內存中的數據顯式地加載到寄存器中。NVCC編譯器將該原語轉換為多個內存加載指令。在運行時runtime，每個線程從矩陣A加載16個元素，從矩陣B加載16個元素。wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag)wmma::load_matrix_sync
準備和算法
將固定大小用于256通道和14 x 14尺寸的輸入張量。批處理大小為256。卷積過濾器包含512個大小為3 x 3的過濾器。對于卷積，使用步幅大小1和填充大小1。在示例中，使用NHWCnc內存布局。以下代碼定義了TVM中的卷積算法。
import tvm
from tvm import te
import numpy as np
from tvm.contrib import nvcc

# The sizes of inputs and filters
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

# Input feature map: (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,
)
# Output feature map: (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”)

# Algorithm
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。在硬件上，所有片段作用域存儲在片上寄存器級別，與本地存儲器位于同一位置。
# Designate the memory hierarchy
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_fragmentload_matrixmma_syncstore_matrixfill_fragmentmma_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,) = outsdef 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變形，并且每個變形調用4x2 TensorCore指令。因此，每個warp的輸出形狀為64x32，每個塊輸出128x128標題。由于共享內存空間的限制，一次只能加載2個塊（2x128x128個圖塊）。
warp操作
請注意，所有TensorCore指令均為warp級指令，這意味著warp中的所有32個線程應同時執行此指令。使theadIdx.x范圍= 32是解決此問題的最簡單方法之一。然后可以將threadIdx.x綁定到任何循環，除了那些直接或間接包含TensorCore內部函數的循環。還要注意，這不是唯一的解決方案。唯一要做的是確保warp中的所有線程可以同時調用TensorCore。
# Define tiling sizes
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)

# Schedule local computation
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)

# Move intermediate computation into each output compute tile
s[AF].compute_at(s[ConvF], kw)
s[WF].compute_at(s[ConvF], kw)

# Schedule for A’s share memory
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)

# Schedule for W’s share memory
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))
輸出：
primfn(A_1: handle, W_1: handle, Conv_1: handle) -> ()
attr = {“global_symbol”: “main”, “tir.noalias”: True}
buffers = {Conv: Buffer(Conv_2: Pointer(float32), float32, [16, 14, 14, 32, 16, 16], []),
W: Buffer(W_2: Pointer(float16), float16, [3, 3, 16, 32, 16, 16], []),
A: Buffer(A_2: Pointer(float16), float16, [16, 14, 14, 16, 16, 16], [])}
buffer_map = {A_1: A, W_1: W, Conv_1: Conv} {
attr [IterVar(blockIdx.z: int32, (nullptr), “ThreadIndex”, “blockIdx.z”)] “thread_extent” = 196;
attr [Conv.wmma.accumulator: Pointer(float32)] “storage_scope” = “wmma.accumulator”;
allocate(Conv.wmma.accumulator, float32, [2048]);
attr [Apad.shared: Pointer(float16)] “storage_scope” = “shared”;
allocate(Apad.shared, float16, [12288]);
attr [W.shared: Pointer(float16)] “storage_scope” = “shared”;
allocate(W.shared, float16, [12288]);
attr [Apad.shared.wmma.matrix_a: Pointer(float16)] “storage_scope” = “wmma.matrix_a”;
allocate(Apad.shared.wmma.matrix_a, float16, [512]);
attr [W.shared.wmma.matrix_b: Pointer(float16)] “storage_scope” = “wmma.matrix_b”;
allocate(W.shared.wmma.matrix_b, float16, [1024]);
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[((((n.c.init1024) + (o.c.init256)) + (nn.c.init16)) + 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) {
attr [IterVar(threadIdx.x: int32, (nullptr), “ThreadIndex”, “threadIdx.x”)] “thread_extent” = 32;
Apad.shared[((((((threadIdx.y3072) + (threadIdx.z1536)) + (ax2512)) + (ax3256)) + (ax4.ax5.fused.outer32)) + 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)), (float16*)A_2[(((((((((((blockIdx.x6422528) + (threadIdx.y1605632)) + (threadIdx.z802816)) + (kh57344)) + (blockIdx.z4096)) + (ax24096)) + (ic.outer512)) + (ax3256)) + (ax4.ax5.fused.outer32)) + 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[ramp((((((ax14096) + (ax2_12048)) + (threadIdx.y512)) + (threadIdx.z256)) + (threadIdx.x8)), 1, 8)] = (float16x8*)W_2[ramp(((((((((kh393216) + (ax1131072)) + (ic.outer16384)) + (ax2_18192)) + (blockIdx.y2048)) + (threadIdx.y512)) + (threadIdx.z256)) + (threadIdx.x8)), 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) {
Apad.shared.wmma.matrix_a[(((ax0256) + (ax416)) + ax5)] = (float16*)Apad.shared[((((((threadIdx.y3072) + (ax01536)) + (kw512)) + (ic.inner256)) + (ax416)) + ax5)]
}
}
}
for (ax3_1: int32, 0, 4) {
for (ax4_1: int32, 0, 16) {
for (ax5_1: int32, 0, 16) {
W.shared.wmma.matrix_b[(((ax3_1256) + (ax4_116)) + ax5_1)] = (float16)W.shared[((((((kw4096) + (ic.inner2048)) + (threadIdx.z1024)) + (ax3_1256)) + (ax4_116)) + 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) {
Conv.wmma.accumulator[((((n.c1024) + (o.c256)) + (nn.c16)) + oo.c)] = ((float32*)Conv.wmma.accumulator[((((n.c1024) + (o.c256)) + (nn.c16)) + oo.c)] + (cast(float32, (float16)Apad.shared.wmma.matrix_a[(((n.c256) + (nn.c16)) + ii)])cast(float32, (float16)W.shared.wmma.matrix_b[(((o.c256) + (ii16)) + oo.c)])))
}
}
}
}
}
}
}
}
}
for (n.inner: int32, 0, 2) {
for (o.inner: int32, 0, 4) {
for (nn: int32, 0, 16) {
for (oo: int32, 0, 16) {
Conv_2[(((((((((blockIdx.x12845056) + (threadIdx.y3211264)) + (n.inner1605632)) + (blockIdx.z8192)) + (blockIdx.y2048)) + (threadIdx.z1024)) + (o.inner256)) + (nn16)) + oo)] = (float32*)Conv.wmma.accumulator[((((n.inner1024) + (o.inner256)) + (nn16)) + oo)]
}
}
}
}
}
}
降低算力
最后一個階段是通過將2D卷積映射到張量特征，來將計算循環降低到TensorCore硬件特征
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))
輸出：
primfn(A_1: handle, W_1: handle, Conv_1: handle) -> ()
attr = {“global_symbol”: “main”, “tir.noalias”: True}
buffers = {Conv: Buffer(Conv_2: Pointer(float32), float32, [16, 14, 14, 32, 16, 16], []),
W: Buffer(W_2: Pointer(float16), float16, [3, 3, 16, 32, 16, 16], []),
A: Buffer(A_2: Pointer(float16), float16, [16, 14, 14, 16, 16, 16], [])}
buffer_map = {A_1: A, W_1: W, Conv_1: Conv} {
attr [IterVar(blockIdx.z: int32, (nullptr), “ThreadIndex”, “blockIdx.z”)] “thread_extent” = 196;
attr [Conv.wmma.accumulator: Pointer(float32)] “storage_scope” = “wmma.accumulator”;
allocate(Conv.wmma.accumulator, float32, [2048]);
attr [Apad.shared: Pointer(float16)] “storage_scope” = “shared”;
allocate(Apad.shared, float16, [12288]);
attr [W.shared: Pointer(float16)] “storage_scope” = “shared”;
allocate(W.shared, float16, [12288]);
attr [Apad.shared.wmma.matrix_a: Pointer(float16)] “storage_scope” = “wmma.matrix_a”;
allocate(Apad.shared.wmma.matrix_a, float16, [512]);
attr [W.shared.wmma.matrix_b: Pointer(float16)] “storage_scope” = “wmma.matrix_b”;
allocate(W.shared.wmma.matrix_b, float16, [1024]);
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.init4) + 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) {
attr [IterVar(threadIdx.x: int32, (nullptr), “ThreadIndex”, “threadIdx.x”)] “thread_extent” = 32;
Apad.shared[((((((threadIdx.y3072) + (threadIdx.z1536)) + (ax2512)) + (ax3256)) + (ax4.ax5.fused.outer32)) + 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)), (float16)A_2[(((((((((((blockIdx.x6422528) + (threadIdx.y1605632)) + (threadIdx.z802816)) + (kh57344)) + (blockIdx.z4096)) + (ax24096)) + (ic.outer512)) + (ax3256)) + (ax4.ax5.fused.outer32)) + 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[ramp((((((ax14096) + (ax2_12048)) + (threadIdx.y512)) + (threadIdx.z256)) + (threadIdx.x8)), 1, 8)] = (float16x8*)W_2[ramp(((((((((kh393216) + (ax1131072)) + (ic.outer16384)) + (ax2_18192)) + (blockIdx.y2048)) + (threadIdx.y512)) + (threadIdx.z256)) + (threadIdx.x8)), 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.y3072) + (ax01536)) + (kw512)) + (ic.inner256)), 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, ((((kw4096) + (ic.inner2048)) + (threadIdx.z1024)) + (ax3_1256)), 256, 1, dtype=handle), 16, “row_major”, dtype=handle)
}
for (n.c: int32, 0, 2) {
for (o.c: int32, 0, 4) {
@tir.tvm_mma_sync(Conv.wmma.accumulator, ((n.c4) + o.c), Apad.shared.wmma.matrix_a, n.c, W.shared.wmma.matrix_b, o.c, Conv.wmma.accumulator, ((n.c4) + o.c), 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.inner4) + o.inner), @tir.tvm_access_ptr(@tir.type_annotation(, dtype=float32), Conv_2, (((((((blockIdx.x12845056) + (threadIdx.y3211264)) + (n.inner1605632)) + (blockIdx.z8192)) + (blockIdx.y2048)) + (threadIdx.z1024)) + (o.inner256)), 256, 2, dtype=handle), 16, “row_major”, dtype=handle)
}
}
}
}
生成CUDA內核
最后，使用TVM生成和編譯CUDA內核，并評估卷積的延遲。由于TensorCores僅在具有Compute Capability 7.0或更高版本的NVIDIA GPU中受支持，因此它可能無法在構建服務器上運行
ctx = tvm.gpu(0)
if nvcc.have_tensorcore(ctx.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, ctx)
w = tvm.nd.array(w_np, ctx)
c = tvm.nd.array(np.zeros(output_shape, dtype=Conv.dtype), ctx)
evaluator = func.time_evaluator(func.entry_name, ctx, number=10)
print(“conv2d with tensor core: %f ms” % (evaluator(a, w, c).mean * 1e3))
輸出：
conv2d with tensor core: 8.329637 ms
概要
本文演示了如何使用TVM調度原語在特定GPU上調用TensorCore。
https://tvm.apache.org/docs/tutorials/optimize/opt_conv_tensorcore.html

總結

以上是生活随笔為你收集整理的如何使用TensorCores优化卷积的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。