张量内积#
%cd ../..
import set_env
/media/pc/data/4tb/lxw/home/lxw/tvm-book/doc/tutorials
import numpy as np
import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
from tvm import te, topi
矩阵乘法分块#
有 \(\mathbf{x} = (x_1, \cdots, x_w)^T, \mathbf{y} = (y_1, \cdots, y_w)^T \in \mathbb{R}^w\)，则它们的内积为
\[
\langle \mathbf{x}, \mathbf{y} \rangle = \sum_i^w x_i y_i = x^T \cdot y = x \cdot y^T \in \mathbb{R}
\]
进一步有 \(\mathbf{X} = (\mathbf{x}_1, \cdots, \mathbf{x}_h)^T \in \mathbb{R}^{h \times w}\), \(\mathbf{Y} = (\mathbf{y}_1, \cdots, \mathbf{y}_{h_o})^T \in \mathbb{R}^{h_o \times w}\)，有
\[
\langle \mathbf{X}, \mathbf{Y} \rangle = \mathbf{X} \cdot \mathbf{Y}^T = (\langle \mathbf{x}_i, \mathbf{y}_j \rangle)_{i=1, j=1}^{i=h, j=h_o} \in \mathbb{R}^{h \times h_o}
\]
a_np = np.arange(24).reshape(3, 8)
b_np = np.arange(16).reshape(2, 8)
print(f"a_np:\n{a_np}\nb_np:\n{b_np}")
a_np:
[[ 0 1 2 3 4 5 6 7]
[ 8 9 10 11 12 13 14 15]
[16 17 18 19 20 21 22 23]]
b_np:
[[ 0 1 2 3 4 5 6 7]
[ 8 9 10 11 12 13 14 15]]
内积参考结果：
c_np = a_np @ b_np.T
c_np
array([[ 140, 364],
[ 364, 1100],
[ 588, 1836]])
tvm 数组：
a_nd = tvm.nd.array(a_np)
b_nd = tvm.nd.array(b_np)
c_nd = tvm.nd.empty(c_np.shape, dtype=c_np.dtype)
m = te.var("m")
n = te.var("n")
o = te.var("o")
A = te.placeholder((m, n), "int64", "X")
B = te.placeholder((o, n), "int64", "Y")
C = topi.matmul(A, B, transp_b=True) # 矩阵乘法
te_func = te.create_prim_func([A, B, C])
te_func.show()
mod = tvm.build(te_func, target="llvm")
mod(a_nd, b_nd, c_nd)
np.testing.assert_equal(c_nd.numpy(), c_np)
# from tvm.script import tir as T
@T.prim_func
def func(var_X: T.handle, var_Y: T.handle, var_T_matmul: T.handle):
# function attr dict
T.func_attr({"global_symbol": "main", "tir.noalias": True})
m = T.var("int32")
n = T.var("int32")
o = T.var("int32")
X = T.match_buffer(var_X, [m, n], dtype="int64")
Y = T.match_buffer(var_Y, [o, n], dtype="int64")
T_matmul = T.match_buffer(var_T_matmul, [m, o], dtype="int64")
# body
# with T.block("root")
for i0, i1, i2 in T.grid(m, o, n):
with T.block("T_matmul"):
ax0, ax1, k = T.axis.remap("SSR", [i0, i1, i2])
T.reads(X[ax0, k], Y[ax1, k])
T.writes(T_matmul[ax0, ax1])
with T.init():
T_matmul[ax0, ax1] = T.int64(0)
T_matmul[ax0, ax1] = T_matmul[ax0, ax1] + X[ax0, k] * Y[ax1, k]
三维张量内积#
对于三维张量 \(\mathsf{X} = (\mathbf{X}_1, \cdots, \mathbf{X}_{c_i})^T \in \mathbb{R}^{c_i \times h \times w}\), \(\mathsf{Y} = (\mathbf{Y}_1, \cdots, \mathbf{Y}_{c_o})^T \in \mathbb{R}^{c_o \times {h_o} \times w}\)，有
\[
\langle \mathsf{X}, \mathsf{Y} \rangle = \mathsf{X} \cdot \mathsf{Y}^T = (\langle \mathbf{X}_i, \mathbf{Y}_j \rangle)_{i=1, j=1}^{i=c_i, j=c_o} \in \mathbb{R}^{c_i \times h \times h_o \times c_o}
\]
mod = IRModule({"mm": te_func})
sch = tvm.tir.Schedule(mod)
block_Z = sch.get_block("T_matmul", func_name="mm")
ax0, ax1, k = sch.get_loops(block_Z)
k0, k1 = sch.split(k, factors=[None, 4])
sch.mod.show()
mod = tvm.build(sch.mod, target="llvm")
mod(a_nd, b_nd, c_nd)
np.testing.assert_equal(c_nd.numpy(), c_np)
# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
@T.prim_func
def mm(var_X: T.handle, var_Y: T.handle, var_T_matmul: T.handle):
# function attr dict
T.func_attr({"global_symbol": "main", "tir.noalias": True})
m = T.var("int32")
n = T.var("int32")
o = T.var("int32")
X = T.match_buffer(var_X, [m, n], dtype="int64")
Y = T.match_buffer(var_Y, [o, n], dtype="int64")
T_matmul = T.match_buffer(var_T_matmul, [m, o], dtype="int64")
# body
# with T.block("root")
for i0, i1, i2_0, i2_1 in T.grid(m, o, (n + 3) // 4, 4):
with T.block("T_matmul"):
T.where(i2_0 * 4 + i2_1 < n)
ax0, ax1 = T.axis.remap("SS", [i0, i1])
k = T.axis.reduce(n, i2_0 * 4 + i2_1)
T.reads(X[ax0, k], Y[ax1, k])
T.writes(T_matmul[ax0, ax1])
with T.init():
T_matmul[ax0, ax1] = T.int64(0)
T_matmul[ax0, ax1] = T_matmul[ax0, ax1] + X[ax0, k] * Y[ax1, k]
四维张量内积#
对于四维张量 \(\mathop{X} = (\mathsf{X}_1, \cdots, \mathsf{X}_{b_i})^T \in \mathbb{R}^{b_i \times c_i \times h \times w}\), \(\mathop{Y} = (\mathsf{Y}_1, \cdots, \mathsf{Y}_{b_o})^T \in \mathbb{R}^{b_o \times c_o \times {h_o} \times w}\)，有
\[
\langle \mathop{X}, \mathop{Y} \rangle = X \cdot Y^T = (\langle \mathsf{X}_i, \mathsf{Y}_j \rangle)_{i=1, j=1}^{i=b_i, j=b_o} \in \mathbb{R}^{b_i \times c_i \times h \times h_o \times c_o \times b_o}
\]
mod = IRModule({"mm": te_func})
sch = tvm.tir.Schedule(mod)
block_Z = sch.get_block("T_matmul", func_name="mm")
ax0, ax1, k = sch.get_loops(block_Z)
k0, k1, k2 = sch.split(k, factors=[None, 2, 2])
sch.mod.show()
mod = tvm.build(sch.mod, target="llvm")
mod(a_nd, b_nd, c_nd)
np.testing.assert_equal(c_nd.numpy(), c_np)
# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
@T.prim_func
def mm(var_X: T.handle, var_Y: T.handle, var_T_matmul: T.handle):
# function attr dict
T.func_attr({"global_symbol": "main", "tir.noalias": True})
m = T.var("int32")
n = T.var("int32")
o = T.var("int32")
X = T.match_buffer(var_X, [m, n], dtype="int64")
Y = T.match_buffer(var_Y, [o, n], dtype="int64")
T_matmul = T.match_buffer(var_T_matmul, [m, o], dtype="int64")
# body
# with T.block("root")
for i0, i1, i2_0, i2_1, i2_2 in T.grid(m, o, (n + 3) // 4, 2, 2):
with T.block("T_matmul"):
T.where((i2_0 * 2 + i2_1) * 2 + i2_2 < n)
ax0, ax1 = T.axis.remap("SS", [i0, i1])
k = T.axis.reduce(n, i2_0 * 4 + i2_1 * 2 + i2_2)
T.reads(X[ax0, k], Y[ax1, k])
T.writes(T_matmul[ax0, ax1])
with T.init():
T_matmul[ax0, ax1] = T.int64(0)
T_matmul[ax0, ax1] = T_matmul[ax0, ax1] + X[ax0, k] * Y[ax1, k]
@tvm.script.ir_module
class MatmulModule:
@T.prim_func
def main(
A: T.Buffer[(1024, 1024), "float32"],
B: T.Buffer[(1024, 1024), "float32"],
C: T.Buffer[(1024, 1024), "float32"],
) -> None:
T.func_attr({"global_symbol": "main", "tir.noalias": True})
for i, j, k in T.grid(1024, 1024, 1024):
with T.block("matmul"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
C[vi, vj] = T.float32(0)
C[vi, vj] += A[vi, vk] * B[vj, vk]
sch = tvm.tir.Schedule(MatmulModule)
i, j, k = sch.get_loops("matmul")
i, ii = sch.split(i, factors=[None, 16])
j, ji = sch.split(j, factors=[None, 16])
k, ki = sch.split(k, factors=[None, 16])
sch.reorder(i, j, k, ii, ji, ki)
sch.mod.show()
# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer[(1024, 1024), "float32"], B: T.Buffer[(1024, 1024), "float32"], C: T.Buffer[(1024, 1024), "float32"]):
# function attr dict
T.func_attr({"tir.noalias": True, "global_symbol": "main"})
# body
# with T.block("root")
for i_0, j_0, k_0, i_1, j_1, k_1 in T.grid(64, 64, 64, 16, 16, 16):
with T.block("matmul"):
vi = T.axis.spatial(1024, i_0 * 16 + i_1)
vj = T.axis.spatial(1024, j_0 * 16 + j_1)
vk = T.axis.reduce(1024, k_0 * 16 + k_1)
T.reads(A[vi, vk], B[vj, vk])
T.writes(C[vi, vj])
with T.init():
C[vi, vj] = T.float32(0)
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
block_mm = sch.blockize(ii)
sch.mod.show()
# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer[(1024, 1024), "float32"], B: T.Buffer[(1024, 1024), "float32"], C: T.Buffer[(1024, 1024), "float32"]):
# function attr dict
T.func_attr({"tir.noalias": True, "global_symbol": "main"})
# body
# with T.block("root")
for i_0, j_0, k_0 in T.grid(64, 64, 64):
with T.block("matmul_o"):
vi_o, vj_o, vk_o = T.axis.remap("SSR", [i_0, j_0, k_0])
T.reads(A[vi_o * 16 : vi_o * 16 + 16, vk_o * 16 : vk_o * 16 + 16], B[vj_o * 16 : vj_o * 16 + 16, vk_o * 16 : vk_o * 16 + 16])
T.writes(C[vi_o * 16 : vi_o * 16 + 16, vj_o * 16 : vj_o * 16 + 16])
with T.init():
for i_1, j_1 in T.grid(16, 16):
with T.block("matmul_init"):
vi_i_init, vj_i_init = T.axis.remap("SS", [i_1, j_1])
T.reads()
T.writes(C[vi_o * 16 + vi_i_init, vj_o * 16 + vj_i_init])
C[vi_o * 16 + vi_i_init, vj_o * 16 + vj_i_init] = T.float32(0)
for i_1, j_1, k_1 in T.grid(16, 16, 16):
with T.block("matmul"):
vi_i, vj_i, vk_i = T.axis.remap("SSR", [i_1, j_1, k_1])
T.reads(C[vi_o * 16 + vi_i, vj_o * 16 + vj_i], A[vi_o * 16 + vi_i, vk_o * 16 + vk_i], B[vj_o * 16 + vj_i, vk_o * 16 + vk_i])
T.writes(C[vi_o * 16 + vi_i, vj_o * 16 + vj_i])
C[vi_o * 16 + vi_i, vj_o * 16 + vj_i] = C[vi_o * 16 + vi_i, vj_o * 16 + vj_i] + A[vi_o * 16 + vi_i, vk_o * 16 + vk_i] * B[vj_o * 16 + vj_i, vk_o * 16 + vk_i]
A_reg = sch.cache_read(block_mm, 0, storage_scope="global.A_reg")
B_reg = sch.cache_read(block_mm, 1, storage_scope="global.B_reg")
sch.compute_at(A_reg, k)
sch.compute_at(B_reg, k)
write_back_block = sch.cache_write(block_mm, 0, storage_scope="global.accumulator")
sch.reverse_compute_at(write_back_block, j)
sch.mod.show()
# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
@T.prim_func
def main(A: T.Buffer[(1024, 1024), "float32"], B: T.Buffer[(1024, 1024), "float32"], C: T.Buffer[(1024, 1024), "float32"]):
# function attr dict
T.func_attr({"tir.noalias": True, "global_symbol": "main"})
# body
# with T.block("root")
A_global_A_reg = T.alloc_buffer([1024, 1024], dtype="float32", scope="global.A_reg")
B_global_B_reg = T.alloc_buffer([1024, 1024], dtype="float32", scope="global.B_reg")
C_global_accumulator = T.alloc_buffer([1024, 1024], dtype="float32", scope="global.accumulator")
for i_0, j_0 in T.grid(64, 64):
for k_0 in T.serial(64):
for ax0, ax1 in T.grid(16, 16):
with T.block("A_global.A_reg"):
v0 = T.axis.spatial(1024, i_0 * 16 + ax0)
v1 = T.axis.spatial(1024, k_0 * 16 + ax1)
T.reads(A[v0, v1])
T.writes(A_global_A_reg[v0, v1])
A_global_A_reg[v0, v1] = A[v0, v1]
for ax0, ax1 in T.grid(16, 16):
with T.block("B_global.B_reg"):
v0 = T.axis.spatial(1024, j_0 * 16 + ax0)
v1 = T.axis.spatial(1024, k_0 * 16 + ax1)
T.reads(B[v0, v1])
T.writes(B_global_B_reg[v0, v1])
B_global_B_reg[v0, v1] = B[v0, v1]
with T.block("matmul_o"):
vi_o, vj_o, vk_o = T.axis.remap("SSR", [i_0, j_0, k_0])
T.reads(A_global_A_reg[vi_o * 16 : vi_o * 16 + 16, vk_o * 16 : vk_o * 16 + 16], B_global_B_reg[vj_o * 16 : vj_o * 16 + 16, vk_o * 16 : vk_o * 16 + 16])
T.writes(C_global_accumulator[vi_o * 16 : vi_o * 16 + 16, vj_o * 16 : vj_o * 16 + 16])
with T.init():
for i_1, j_1 in T.grid(16, 16):
with T.block("matmul_init"):
vi_i_init, vj_i_init = T.axis.remap("SS", [i_1, j_1])
T.reads()
T.writes(C_global_accumulator[vi_o * 16 + vi_i_init, vj_o * 16 + vj_i_init])
C_global_accumulator[vi_o * 16 + vi_i_init, vj_o * 16 + vj_i_init] = T.float32(0)
for i_1, j_1, k_1 in T.grid(16, 16, 16):
with T.block("matmul"):
vi_i, vj_i, vk_i = T.axis.remap("SSR", [i_1, j_1, k_1])
T.reads(C_global_accumulator[vi_o * 16 + vi_i, vj_o * 16 + vj_i], A_global_A_reg[vi_o * 16 + vi_i, vk_o * 16 + vk_i], B_global_B_reg[vj_o * 16 + vj_i, vk_o * 16 + vk_i])
T.writes(C_global_accumulator[vi_o * 16 + vi_i, vj_o * 16 + vj_i])
C_global_accumulator[vi_o * 16 + vi_i, vj_o * 16 + vj_i] = C_global_accumulator[vi_o * 16 + vi_i, vj_o * 16 + vj_i] + A_global_A_reg[vi_o * 16 + vi_i, vk_o * 16 + vk_i] * B_global_B_reg[vj_o * 16 + vj_i, vk_o * 16 + vk_i]
for ax0, ax1 in T.grid(16, 16):
with T.block("C_global.accumulator"):
v0 = T.axis.spatial(1024, i_0 * 16 + ax0)
v1 = T.axis.spatial(1024, j_0 * 16 + ax1)
T.reads(C_global_accumulator[v0, v1])
T.writes(C[v0, v1])
C[v0, v1] = C_global_accumulator[v0, v1]