exercises有答案吗？ #34

hbsun2113 · 2022-07-06T12:15:59Z

hbsun2113
Jul 6, 2022

请问exercises会公布答案吗？
我写了一版供大家参考。最后一个练习的最后一个步骤出错了，请大家帮忙看看问题在哪里？
2.5.1.2. Exercise 1: Broadcast Add：

@tvm.script.ir_module
class MyAdd:
  @T.prim_func
  def add(A: T.Buffer[(4, 4), "int64"],
          B: T.Buffer[(4), "int64"],
          C: T.Buffer[(4, 4), "int64"]):
    T.func_attr({"global_symbol": "add", "tir.noalias": True})
    for i, j in T.grid(4, 4):
      with T.block("C"):
        vi=T.axis.spatial(4, i)
        vj=T.axis.spatial(4, j)
        C[vi, vj] = A[vi, vj] + B[vj]

2.5.1.3. Exercise 2: 2D Convolution

# hbsun: canot use variable H and W: https://github.com/mlc-ai/mlc-zh/discussions/33
N, CI, H, W, CO, K = 1, 1, 8, 8, 2, 3
OUT_H, OUT_W = H - K + 1, W - K + 1
data = np.arange(N*CI*H*W).reshape(N, CI, H, W)
weight = np.arange(CO*CI*K*K).reshape(CO, CI, K, K)

@tvm.script.ir_module
class MyConv:
  @T.prim_func
  def conv(A: T.Buffer[(N, CI, 8, 8), "int64"], B: T.Buffer[(CO, CI, K, K), "int64"], C: T.Buffer[(N, CO, OUT_H, OUT_W), "int64"]):
    T.func_attr({"global_symbol": "conv", "tir.noalias": True})
    for b, k, i, j, q, di, dj in T.grid(N, CO, OUT_H, OUT_W, CI, K, K):
      with T.block("C"):
        vb=T.axis.spatial(N, b)
        vk=T.axis.spatial(CO, k)
        vi=T.axis.spatial(OUT_H, i)
        vj=T.axis.spatial(OUT_W, j)
        vq=T.axis.reduce(CI, q)
        vdi=T.axis.reduce(K, di)
        vdj=T.axis.reduce(K, dj)
        
        # hbsun: need init first
        with T.init():
          C[vb, vk, vi, vj] = T.int64(0)
        C[vb, vk, vi, vj] += A[vb, vq, vi+vdi, vj+vdj] * B[vk, vq, vdi, vdj]

2.5.2.2. Exercise 3: Transform a batch matmul program
3.1 bmm_relu

@tvm.script.ir_module
class MyBmmRule:
  @T.prim_func
  def bmm_relu(A: T.Buffer[(16, 128, 128), "float32"], B: T.Buffer[(16, 128, 128), "float32"], C: T.Buffer[(16, 128, 128), "float32"]):
    T.func_attr({"global_symbol": "bmm_relu", "tir.noalias": True})
    Y = T.alloc_buffer([16, 128, 128], dtype="float32")
    for n, i, j, k in T.grid(16, 128, 128, 128):
      with T.block("Y"):
        vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k])
        with T.init():
          Y[vn, vi, vj]=T.float32(0)
        Y[vn, vi, vj] += A[vn, vi, vk] * B[vn, vk, vj]
    for n, i, j in T.grid(16, 128, 128):
      with T.block("C"):
        vn, vi, vj = T.axis.remap("SSS", [n, i, j])
        C[vn, vi, vj] = T.max(Y[vn, vi, vj], T.float32(0))

sch = tvm.tir.Schedule(MyBmmRule)
print(sch.mod.script())

3.2 Also please validate your result

# hbsun: please note the official code has an error: https://github.com/mlc-ai/mlc-zh/discussions/29
dtype = "float32"
X3, Y3, Z3 = 16, 128, 128
A3 = np.arange(X3*Y3*Z3).astype(dtype).reshape(X3, Y3, Z3)
B3 = np.random.rand(X3*Y3*Z3).astype(dtype).reshape(X3, Y3, Z3)
C3 = np.empty((X3, Y3, Z3), dtype=dtype)
lnumpy_mm_relu_v2(A3, B3, C3)
print(C3)

bmm_relu_lib = tvm.build(MyBmmRule, target="llvm")
bmm_relu_A3 = tvm.nd.array(A3)
bmm_relu_B3 = tvm.nd.array(B3)
bmm_relu_C3 = tvm.nd.array(np.empty((X3, Y3, Z3), dtype=dtype))
bmm_relu_lib["bmm_relu"](bmm_relu_A3, bmm_relu_B3, bmm_relu_C3)
np.testing.assert_allclose(bmm_relu_C3.numpy(), C3, rtol=1e-5)

# 或者使用官方code:
before_rt_lib = tvm.build(MyBmmRule, target="llvm")
a_numpy = np.random.rand(16, 128, 128).astype("float32")
b_numpy = np.random.rand(16, 128, 128).astype("float32")
c_numpy = np.random.rand(16, 128, 128).astype("float32")
a_tvm = tvm.nd.array(a_numpy)
b_tvm = tvm.nd.array(b_numpy)
c_tvm = tvm.nd.array(c_numpy)
lnumpy_mm_relu_v2(a_numpy, b_numpy, c_numpy)
bmm_relu_lib["bmm_relu"](a_tvm, b_tvm, c_tvm)
np.testing.assert_allclose(c_tvm.numpy(), c_numpy, rtol=1e-5)

3.3 Step 1-4

# TODO: transformations
# Hints: you can use
# `IPython.display.Code(sch.mod.script(), language="python")`
# or `print(sch.mod.script())`
# to show the current program at any time during the transformation.
sch = tvm.tir.Schedule(MyBmmRule)
# Step 1. Get blocks
block_Y = sch.get_block("Y", func_name="bmm_relu")

# Step 2. Get loops
b, i, j, k = sch.get_loops(block_Y)

# Step 3. Organize the loops
k0, k1 = sch.split(k, factors=[32, 4])
j0, j1 = sch.split(j, factors=[16, 8])
                  
sch.reorder(j0, k0, k1, j1)
print(sch.mod.script())

block_C = sch.get_block("C", func_name="bmm_relu")
print(block_C)
print(k0)
sch.reverse_compute_at(block_C, j0)
print(sch.mod.script())

# Step 4. decompose reduction
Y_init = sch.decompose_reduction(block_Y, k0)
print(sch.mod.script())

3.4 Step 5: vectorize / unroll

# Step 5. vectorize / parallel / unroll
n, i, j_0, j_1_init = sch.get_loops(Y_init)
sch.vectorize(j_1_init)

n, i, j_0, i2_1 = sch.get_loops(block_C)
sch.vectorize(i2_1)
print(sch.mod.script())


block_Y_update = sch.get_block("Y_update", func_name="bmm_relu")
n, i, j_0, k_0, k_1, j_1 = sch.get_loops(block_Y_update)
sch.unroll(k_1)
print(sch.mod.script())

3.5 Step 5: parallel: there is an error

n, i, j_0, k_0, k_1, j_1 = sch.get_loops(block_Y_update)
sch.parallel(n)

jzh18 · 2022-07-06T19:02:17Z

jzh18
Jul 6, 2022

最后一题我是这么写的：

sch = tvm.tir.Schedule(MyBmmRule)
# TODO: transformations
# Hints: you can use
# `IPython.display.Code(sch.mod.script(), language="python")`
# or `print(sch.mod.script())`
# to show the current program at any time during the transformation.

# Step 1. Get blocks
Y = sch.get_block("Y", func_name="bmm_relu")


# Step 2. Get loops
b, i, j, k = sch.get_loops(Y)
sch.parallel(b)

# Step 3. Organize the loops
j0, j1 = sch.split(j, factors=[None, 8])
sch.reorder(j0,k, j1)
block_C = sch.get_block("C", func_name="bmm_relu")
sch.reverse_compute_at(block_C, j0)

# Step 4. decompose reduction
Y_init = sch.decompose_reduction(Y, k)
n, i, j_0, j_1_init = sch.get_loops(Y_init)
_, _, _, ax0 = sch.get_loops(block_C)
Y_update_block = sch.get_block("Y_update", func_name="bmm_relu")
_, _, _, k, j_1 = sch.get_loops(Y_update_block)
k0, k1 = sch.split(k, factors=[32, 4])


# # Step 5. vectorize / parallel / unroll
sch.vectorize(j_1_init)
sch.vectorize(ax0)
sch.unroll(k1)

print(sch.mod.script())

通过了测试，完整代码可以参考：https://github.com/JZh18/mlc-ai-exercises/blob/main/2_5_Exercises_for_TensorIR.ipynb