def strided_head(H, s_h, s_w):
n, c1, h, w, c0 = H.shape
out_shape = (n, c1, (h - 1) * s_h + 1, (w - 1) * s_w + 1, c0)
H_strided = akg.tvm.compute(out_shape, lambda i0, i1, i2, i3, i4:
akg.tvm.expr.Select(akg.tvm.any(truncmod(i2, s_h) != 0,
truncmod(i3, s_w) != 0),
akg.tvm.const(0.0, dtype="float16"),
H[i0, i1, floordiv(i2, s_h), floordiv(i3, s_w), i4]),
name=H.name + "_strided")
return H_strided
def flip_weight(B, k_c, k_hw, const_shift):
out_shape = (B.shape[1].value * k_hw, k_c // block_size, block_size, block_size)
B_flip = akg.tvm.compute(out_shape,
lambda i0, i1, i2, i3:
B[i1 * k_hw + const_shift - truncmod(i0, k_hw), floordiv(i0, k_hw), i3, i2],
name=B.name + "_flipped")
return B_flip
def transpose_data(A):
out_shape = (A.shape[1] * block_size, truncdiv(A.shape[0], block_size),
A.shape[2], A.shape[3], block_size)
A_transpose = akg.tvm.compute(
out_shape,
lambda j0, j1, j2, j3, j4: A[j1 * block_size + j4,
truncdiv(j0, block_size), j2, j3,
truncmod(j0, block_size)],
name=A.name + "_transposed")
return A_transpose
