wehub-resource-sync
356 lines
14 KiB
Python
356 lines
14 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=missing-function-docstring,missing-module-docstring
|
|
# ruff: noqa: F401
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import tirx
|
|
from tvm.s_tir.schedule.schedule import ScheduleError
|
|
from tvm.s_tir.schedule.testing import (
|
|
assert_structural_equal_ignore_global_symbol,
|
|
verify_trace_roundtrip,
|
|
)
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def transpose_elementwise(
|
|
A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")
|
|
) -> None:
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("B"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vj, vi] * 2.0
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def transpose_elementwise_reindex_read(
|
|
A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")
|
|
) -> None:
|
|
A_reindex = T.sblock_alloc_buffer((128, 128), "float32")
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("A_reindex"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
A_reindex[vi, vj] = A[vj, vi]
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("B"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A_reindex[vi, vj] * 2.0
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def conv2d_nhwc(
|
|
Input: T.Buffer((1, 224, 224, 3), "float32"),
|
|
Weight: T.Buffer((7, 7, 3, 64), "float32"),
|
|
Conv2d_nhwc: T.Buffer((1, 112, 112, 64), "float32"),
|
|
) -> None:
|
|
PadInput = T.sblock_alloc_buffer([1, 230, 230, 3], dtype="float32")
|
|
for i0, i1, i2, i3 in T.grid(1, 230, 230, 3):
|
|
with T.sblock("PadInput"):
|
|
i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
PadInput[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(
|
|
((((i1_1 >= 3) and (i1_1 < 227)) and (i2_1 >= 3)) and (i2_1 < 227)),
|
|
Input[i0_1, (i1_1 - 3), (i2_1 - 3), i3_1],
|
|
T.float32(0),
|
|
dtype="float32",
|
|
)
|
|
for i0, i1, i2, i3, i4, i5, i6 in T.grid(1, 112, 112, 64, 7, 7, 3):
|
|
with T.sblock("conv2d_nhwc"):
|
|
n, h, w, co, rh, rw, rc = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
|
|
with T.init():
|
|
Conv2d_nhwc[n, h, w, co] = T.float32(0)
|
|
Conv2d_nhwc[n, h, w, co] = Conv2d_nhwc[n, h, w, co] + (
|
|
PadInput[n, ((h * 2) + rh), ((w * 2) + rw), ((T.floordiv(co, 64) * 3) + rc)]
|
|
* Weight[rh, rw, rc, co]
|
|
)
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def conv2d_nhwc_reindex_data(
|
|
Input: T.Buffer((1, 224, 224, 3), "float32"),
|
|
Weight: T.Buffer((7, 7, 3, 64), "float32"),
|
|
Conv2d_nhwc: T.Buffer((1, 112, 112, 64), "float32"),
|
|
) -> None:
|
|
PadInput = T.sblock_alloc_buffer([1, 230, 230, 3], dtype="float32")
|
|
ReindexInput = T.sblock_alloc_buffer([1, 112, 112, 7, 7, 3], dtype="float32")
|
|
for i0, i1, i2, i3 in T.grid(1, 230, 230, 3):
|
|
with T.sblock("PadInput"):
|
|
i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
PadInput[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(
|
|
((((i1_1 >= 3) and (i1_1 < 227)) and (i2_1 >= 3)) and (i2_1 < 227)),
|
|
Input[i0_1, (i1_1 - 3), (i2_1 - 3), i3_1],
|
|
T.float32(0),
|
|
dtype="float32",
|
|
)
|
|
for i0, i1, i2, i3, i4, i5 in T.grid(1, 112, 112, 7, 7, 3):
|
|
with T.sblock("ReindexInput"):
|
|
n, h, w, rh, rw, rc = T.axis.remap("SSSSSS", [i0, i1, i2, i3, i4, i5])
|
|
ReindexInput[n, h, w, rh, rw, rc] = PadInput[n, ((h * 2) + rh), ((w * 2) + rw), rc]
|
|
for i0, i1, i2, i3, i4, i5, i6 in T.grid(1, 112, 112, 64, 7, 7, 3):
|
|
with T.sblock("conv2d_nhwc"):
|
|
n, h, w, co, rh, rw, rc = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
|
|
with T.init():
|
|
Conv2d_nhwc[n, h, w, co] = T.float32(0)
|
|
Conv2d_nhwc[n, h, w, co] = Conv2d_nhwc[n, h, w, co] + (
|
|
ReindexInput[n, h, w, rh, rw, rc] * Weight[rh, rw, rc, co]
|
|
)
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def conv2d_nhwc_reindex_weight(
|
|
var_inputs: T.handle, var_weight: T.handle, var_conv2d_nhwc: T.handle
|
|
) -> None:
|
|
inputs = T.match_buffer(var_inputs, [1, 224, 224, 3], dtype="float32")
|
|
weight = T.match_buffer(var_weight, [7, 7, 3, 64], dtype="float32")
|
|
conv2d_nhwc = T.match_buffer(var_conv2d_nhwc, [1, 112, 112, 64], dtype="float32")
|
|
PadInput = T.sblock_alloc_buffer([1, 230, 230, 3], dtype="float32")
|
|
weight_reindex = T.sblock_alloc_buffer([64, 7, 7, 3], dtype="float32")
|
|
for i0, i1, i2, i3 in T.grid(1, 230, 230, 3):
|
|
with T.sblock("PadInput"):
|
|
i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
|
T.reads(inputs[i0_1, i1_1 - 3, i2_1 - 3, i3_1])
|
|
T.writes(PadInput[i0_1, i1_1, i2_1, i3_1])
|
|
PadInput[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(
|
|
i1_1 >= 3 and i1_1 < 227 and i2_1 >= 3 and i2_1 < 227,
|
|
inputs[i0_1, i1_1 - 3, i2_1 - 3, i3_1],
|
|
T.float32(0),
|
|
dtype="float32",
|
|
)
|
|
for ax3, ax4, ax5, ax6 in T.grid(64, 7, 7, 3):
|
|
with T.sblock("weight_reindex"):
|
|
v3, v4, v5, v6 = T.axis.remap("SSSS", [ax3, ax4, ax5, ax6])
|
|
T.reads(weight[v4, v5, v6, v3])
|
|
T.writes(weight_reindex[v3, v4, v5, v6])
|
|
weight_reindex[v3, v4, v5, v6] = weight[v4, v5, v6, v3]
|
|
for i0, i1, i2, i3, i4, i5, i6 in T.grid(1, 112, 112, 64, 7, 7, 3):
|
|
with T.sblock("conv2d_nhwc"):
|
|
n, h, w, co, rh, rw, rc = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
|
|
T.reads(
|
|
PadInput[n, h * 2 + rh, w * 2 + rw, co // 64 * 3 + rc],
|
|
weight_reindex[co, rh, rw, rc],
|
|
)
|
|
T.writes(conv2d_nhwc[n, h, w, co])
|
|
with T.init():
|
|
conv2d_nhwc[n, h, w, co] = T.float32(0)
|
|
conv2d_nhwc[n, h, w, co] = (
|
|
conv2d_nhwc[n, h, w, co]
|
|
+ PadInput[n, h * 2 + rh, w * 2 + rw, co // 64 * 3 + rc]
|
|
* weight_reindex[co, rh, rw, rc]
|
|
)
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def matmul(
|
|
A: T.Buffer((512, 512), "float32"),
|
|
B: T.Buffer((512, 512), "float32"),
|
|
C: T.Buffer((512, 512), "float32"),
|
|
) -> None:
|
|
for i0, i1, i2 in T.grid(512, 512, 512):
|
|
with T.sblock("matmul"):
|
|
i, j, k = T.axis.remap("SSR", [i0, i1, i2])
|
|
T.reads(C[i, j], A[i, k], B[k, j])
|
|
T.writes(C[i, j])
|
|
with T.init():
|
|
C[i, j] = T.float32(0)
|
|
C[i, j] = C[i, j] + A[i, k] * B[k, j]
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def matmul_reindex_write(
|
|
A: T.Buffer((512, 512), "float32"),
|
|
B: T.Buffer((512, 512), "float32"),
|
|
C: T.Buffer((512, 512), "float32"),
|
|
) -> None:
|
|
C_reindex = T.sblock_alloc_buffer([512, 512], dtype="float32")
|
|
for i0, i1, i2 in T.grid(512, 512, 512):
|
|
with T.sblock("matmul"):
|
|
i, j, k = T.axis.remap("SSR", [i0, i1, i2])
|
|
T.reads(C_reindex[i, j], A[i, k], B[k, j])
|
|
T.writes(C_reindex[i, j])
|
|
with T.init():
|
|
C_reindex[i, j] = T.float32(0)
|
|
C_reindex[i, j] = C_reindex[i, j] + A[i, k] * B[k, j]
|
|
for i0, i1 in T.grid(512, 512):
|
|
with T.sblock("C_reindex"):
|
|
v0, v1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(C_reindex[v0, v1])
|
|
T.writes(C[v0, v1])
|
|
C[v0, v1] = C_reindex[v0, v1]
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def multiple_read(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32")) -> None:
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("B"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vj, vi] + A[vi, vj]
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def mixed_dtype(
|
|
p0: T.Buffer((T.int64(2), 1280), "float16"),
|
|
p1: T.Buffer((1280, 1280), "float16"),
|
|
T_matmul_NT: T.Buffer((T.int64(2), 1280), "float16"),
|
|
) -> None:
|
|
for i0, i1, i2 in T.grid(T.int64(2), 1280, 1280):
|
|
with T.sblock("T_matmul_NT"):
|
|
i = T.axis.spatial(T.int64(2), i0)
|
|
j, k = T.axis.remap("SR", [i1, i2])
|
|
T.reads(p0[i, k], p1[j, k])
|
|
T.writes(T_matmul_NT[i, j])
|
|
with T.init():
|
|
T_matmul_NT[i, j] = T.float16(0)
|
|
T_matmul_NT[i, j] = T_matmul_NT[i, j] + p0[i, k] * p1[j, k]
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def mixed_dtype_reindex_write(
|
|
p0: T.Buffer((T.int64(2), 1280), "float16"),
|
|
p1: T.Buffer((1280, 1280), "float16"),
|
|
T_matmul_NT: T.Buffer((T.int64(2), 1280), "float16"),
|
|
) -> None:
|
|
T_matmul_NT_reindex = T.sblock_alloc_buffer([T.int64(2), 1280], dtype="float16")
|
|
for i0, i1, i2 in T.grid(T.int64(2), 1280, 1280):
|
|
with T.sblock("T_matmul_NT"):
|
|
i = T.axis.spatial(T.int64(2), i0)
|
|
j, k = T.axis.remap("SR", [i1, i2])
|
|
T.reads(p0[i, k], p1[j, k])
|
|
T.writes(T_matmul_NT_reindex[i, j])
|
|
with T.init():
|
|
T_matmul_NT_reindex[i, j] = T.float16(0)
|
|
T_matmul_NT_reindex[i, j] = T_matmul_NT_reindex[i, j] + p0[i, k] * p1[j, k]
|
|
for ax0, ax1 in T.grid(T.int64(2), 1280):
|
|
with T.sblock("T_matmul_NT_reindex"):
|
|
v0 = T.axis.spatial(T.int64(2), ax0)
|
|
v1 = T.axis.remap("S", [ax1])
|
|
T.reads(T_matmul_NT_reindex[v0, v1])
|
|
T.writes(T_matmul_NT[v0, v1])
|
|
T_matmul_NT[v0, v1] = T_matmul_NT_reindex[v0, v1]
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def matmul_unit_dim(
|
|
A: T.Buffer((1, 512), "float32"),
|
|
B: T.Buffer((512, 1), "float32"),
|
|
C: T.Buffer((1, 1), "float32"),
|
|
) -> None:
|
|
for i0, i1, i2 in T.grid(1, 1, 512):
|
|
with T.sblock("matmul"):
|
|
i, j, k = T.axis.remap("SSR", [i0, i1, i2])
|
|
T.reads(C[i, j], A[i, k], B[k, j])
|
|
T.writes(C[i, j])
|
|
with T.init():
|
|
C[i, j] = T.float32(0)
|
|
C[i, j] = C[i, j] + A[i, k] * B[k, j]
|
|
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def matmul_unit_dim_reindex_write(
|
|
A: T.Buffer((1, 512), "float32"),
|
|
B: T.Buffer((512, 1), "float32"),
|
|
C: T.Buffer((1, 1), "float32"),
|
|
) -> None:
|
|
C_reindex = T.sblock_alloc_buffer([1, 1], dtype="float32")
|
|
for i0, i1, i2 in T.grid(1, 1, 512):
|
|
with T.sblock("matmul"):
|
|
i, j, k = T.axis.remap("SSR", [i0, i1, i2])
|
|
T.reads(C_reindex[i, j], A[i, k], B[k, j])
|
|
T.writes(C_reindex[i, j])
|
|
with T.init():
|
|
C_reindex[i, j] = T.float32(0)
|
|
C_reindex[i, j] = C_reindex[i, j] + A[i, k] * B[k, j]
|
|
for i0, i1 in T.grid(1, 1):
|
|
with T.sblock("C_reindex"):
|
|
v0, v1 = T.axis.remap("SS", [i0, i1])
|
|
T.reads(C_reindex[v0, v1])
|
|
T.writes(C[v0, v1])
|
|
C[v0, v1] = C_reindex[v0, v1]
|
|
|
|
|
|
use_block_name = tvm.testing.parameter(by_dict={"block_obj": False, "block_name": True})
|
|
use_buffer_name = tvm.testing.parameter(by_dict={"buffer_index": False, "buffer_name": True})
|
|
|
|
|
|
def test_reindex_read_basic(use_block_name, use_buffer_name):
|
|
sch = tvm.s_tir.Schedule(transpose_elementwise)
|
|
block = "B" if use_block_name else sch.get_sblock("B")
|
|
buf = "A" if use_buffer_name else ("read", 0)
|
|
sch.reindex(block, buf)
|
|
assert_structural_equal_ignore_global_symbol(
|
|
transpose_elementwise_reindex_read, sch.mod["main"]
|
|
)
|
|
verify_trace_roundtrip(sch=sch, mod=transpose_elementwise)
|
|
|
|
|
|
def test_conv2d_reindex_weight(use_block_name, use_buffer_name):
|
|
sch = tvm.s_tir.Schedule(conv2d_nhwc)
|
|
block = "conv2d_nhwc" if use_block_name else sch.get_sblock("conv2d_nhwc")
|
|
buf = "Weight" if use_buffer_name else ("read", 1)
|
|
sch.reindex(block, buf)
|
|
assert_structural_equal_ignore_global_symbol(conv2d_nhwc_reindex_weight, sch.mod["main"])
|
|
verify_trace_roundtrip(sch=sch, mod=conv2d_nhwc)
|
|
|
|
|
|
def test_conv2d_reindex_data(use_block_name, use_buffer_name):
|
|
sch = tvm.s_tir.Schedule(conv2d_nhwc)
|
|
block = "conv2d_nhwc" if use_block_name else sch.get_sblock("conv2d_nhwc")
|
|
buf = "PadInput" if use_buffer_name else ("read", 0)
|
|
sch.reindex(block, buf)
|
|
assert_structural_equal_ignore_global_symbol(conv2d_nhwc_reindex_data, sch.mod["main"])
|
|
verify_trace_roundtrip(sch=sch, mod=conv2d_nhwc)
|
|
|
|
|
|
def test_matmul_reindex_write(use_block_name, use_buffer_name):
|
|
sch = tvm.s_tir.Schedule(matmul)
|
|
block = "matmul" if use_block_name else sch.get_sblock("matmul")
|
|
buf = "C" if use_buffer_name else ("write", 0)
|
|
sch.reindex(block, buf)
|
|
assert_structural_equal_ignore_global_symbol(matmul_reindex_write, sch.mod["main"])
|
|
verify_trace_roundtrip(sch=sch, mod=matmul)
|
|
|
|
|
|
def test_reindex_fail_multiple_read(use_block_name, use_buffer_name):
|
|
sch = tvm.s_tir.Schedule(multiple_read)
|
|
block = "B" if use_block_name else sch.get_sblock("B")
|
|
buf = "A" if use_buffer_name else ("read", 0)
|
|
with pytest.raises(ScheduleError):
|
|
sch.reindex(block, buf)
|
|
|
|
|
|
def test_reindex_mixed_dtype(use_block_name, use_buffer_name):
|
|
sch = tvm.s_tir.Schedule(mixed_dtype)
|
|
block = "T_matmul_NT" if use_block_name else sch.get_sblock("T_matmul_NT")
|
|
buf = "T_matmul_NT" if use_buffer_name else ("write", 0)
|
|
sch.reindex(block, buf)
|
|
assert_structural_equal_ignore_global_symbol(mixed_dtype_reindex_write, sch.mod["main"])
|
|
verify_trace_roundtrip(sch=sch, mod=mixed_dtype)
|
|
|
|
|
|
def test_matmul_unit_dim_reindex_write(use_block_name, use_buffer_name):
|
|
sch = tvm.s_tir.Schedule(matmul_unit_dim)
|
|
block = "matmul" if use_block_name else sch.get_sblock("matmul")
|
|
buf = "C" if use_buffer_name else ("write", 0)
|
|
sch.reindex(block, buf)
|
|
assert_structural_equal_ignore_global_symbol(matmul_unit_dim_reindex_write, sch.mod["main"])
|
|
verify_trace_roundtrip(sch=sch, mod=matmul_unit_dim)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|