wehub-resource-sync
460 lines
15 KiB
Python
460 lines
15 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-docstring
|
|
# ruff: noqa: F401, F841
|
|
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.s_tir import Schedule
|
|
from tvm.s_tir.meta_schedule.testing import te_workload
|
|
from tvm.s_tir.schedule.analysis import (
|
|
TensorizeInfo,
|
|
get_auto_tensorize_mapping_info,
|
|
get_tensorize_loop_mapping,
|
|
is_output_block,
|
|
suggest_index_map,
|
|
)
|
|
from tvm.s_tir.tensor_intrin.cuda import (
|
|
WMMA_SYNC_16x16x16_f16f16f16_INTRIN,
|
|
WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
|
|
)
|
|
from tvm.s_tir.tensor_intrin.x86 import dot_product_16x4_u8i8i32_desc
|
|
from tvm.script import tirx as T
|
|
from tvm.te import create_prim_func
|
|
from tvm.tirx import (
|
|
Evaluate,
|
|
For,
|
|
ForKind,
|
|
IndexMap,
|
|
Var,
|
|
decl_buffer,
|
|
floordiv,
|
|
floormod,
|
|
)
|
|
from tvm.tirx.analysis import expr_deep_equal
|
|
from tvm.tirx.function import TensorIntrin
|
|
from tvm.tirx.stmt_functor import pre_order_visit
|
|
|
|
|
|
def _make_vars(*args: str) -> list[Var]:
|
|
return [Var(arg, dtype="int32") for arg in args]
|
|
|
|
|
|
def _make_loops(loop_vars: list[Var], extents: list[int]) -> list[For]:
|
|
assert len(loop_vars) == len(extents)
|
|
return [
|
|
For(
|
|
loop_var=loop_var,
|
|
min=0,
|
|
extent=extent,
|
|
kind=ForKind.SERIAL,
|
|
body=Evaluate(0),
|
|
)
|
|
for loop_var, extent in zip(loop_vars, extents)
|
|
]
|
|
|
|
|
|
def test_suggest_index_map_simple():
|
|
i, j = _make_vars("i", "j")
|
|
index_map = suggest_index_map(
|
|
buffer=decl_buffer(shape=[8, 256]),
|
|
indices=[
|
|
floordiv(i, 16) * 4 + floordiv(j, 16),
|
|
floormod(i, 16) * 16 + floormod(j, 16),
|
|
],
|
|
loops=_make_loops(
|
|
loop_vars=[i, j],
|
|
extents=[32, 64],
|
|
),
|
|
predicate=True,
|
|
)
|
|
expected_index_map = IndexMap.from_func(
|
|
lambda x, y: [
|
|
floordiv(x, 4),
|
|
floordiv(y, 16),
|
|
floormod(x, 4),
|
|
floormod(y, 16),
|
|
],
|
|
)
|
|
assert index_map.is_equivalent_to(expected_index_map)
|
|
|
|
|
|
def test_suggest_index_map_bijective():
|
|
i, j = _make_vars("i", "j")
|
|
index_map = suggest_index_map(
|
|
buffer=decl_buffer(shape=[8]),
|
|
indices=[floormod(j, 4) * 2 + i],
|
|
loops=_make_loops(
|
|
loop_vars=[i, j],
|
|
extents=[2, 32],
|
|
),
|
|
predicate=True,
|
|
)
|
|
expected_index_map = IndexMap.from_func(
|
|
lambda x: [
|
|
floormod(x, 2),
|
|
floordiv(x, 2),
|
|
],
|
|
)
|
|
assert index_map.is_equivalent_to(expected_index_map)
|
|
|
|
|
|
def test_suggest_index_map_winograd():
|
|
"""use case in winograd conv where the indices are complicated"""
|
|
fused_outer, i3_3_fused, i4_0, i4_1 = _make_vars("fused_outer", "i3_3_fused", "i4_0", "i4_1")
|
|
eps = floordiv(fused_outer, 336) * 2 + floordiv(floormod(fused_outer, 16), 8)
|
|
nu = floordiv(floormod(fused_outer, 336), 112) * 2 + floordiv(floormod(fused_outer, 8), 4)
|
|
co = floormod(fused_outer, 4) * 32 + i3_3_fused
|
|
ci = (i4_0 * 32) + i4_1
|
|
buffer = decl_buffer(shape=[6, 6, 128, 128])
|
|
index_map = suggest_index_map(
|
|
buffer=buffer,
|
|
indices=[eps, nu, co, ci],
|
|
loops=_make_loops(
|
|
loop_vars=[fused_outer, i3_3_fused, i4_0, i4_1],
|
|
extents=[1008, 32, 4, 32],
|
|
),
|
|
predicate=True,
|
|
)
|
|
expected_index_map = IndexMap.from_func(
|
|
lambda i0, i1, i2, i3: (
|
|
floordiv(i0, 2),
|
|
floordiv(i1, 2),
|
|
floormod(i0, 2),
|
|
floormod(i1, 2) * 4 + floordiv(i2, 32),
|
|
floormod(i2, 32),
|
|
floordiv(i3, 32),
|
|
floormod(i3, 32),
|
|
)
|
|
)
|
|
assert index_map.is_equivalent_to(expected_index_map)
|
|
inverse_index_map = index_map.inverse(buffer.shape)
|
|
expected_inverse_index_map = IndexMap.from_func(
|
|
lambda i0, i1, i2, i3, i4, i5, i6: (
|
|
((i0 * 2) + i2),
|
|
i1 * 2 + floordiv(i3, 4),
|
|
floormod(i3, 4) * 32 + i4,
|
|
((i5 * 32) + i6),
|
|
)
|
|
)
|
|
assert inverse_index_map.is_equivalent_to(expected_inverse_index_map)
|
|
|
|
|
|
@tvm.script.ir_module
|
|
class DenseTIRModule:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
placeholder: T.Buffer((1024, 1024), "uint8"),
|
|
placeholder_1: T.Buffer((64, 256, 16, 4), "int8"),
|
|
compute: T.Buffer((1024, 1024), "int32"),
|
|
) -> None:
|
|
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
|
with T.sblock("root"):
|
|
T.reads()
|
|
T.writes()
|
|
for i0, i1, i2 in T.grid(1024, 1024, 1024):
|
|
with T.sblock("compute"):
|
|
i, j, k = T.axis.remap("SSR", [i0, i1, i2])
|
|
T.reads(placeholder[i, k], placeholder_1[j // 16, k // 4, j % 16, k % 4])
|
|
T.writes(compute[i, j])
|
|
with T.init():
|
|
compute[i, j] = 0
|
|
compute[i, j] = compute[i, j] + T.cast(placeholder[i, k], "int32") * T.cast(
|
|
placeholder_1[j // 16, k // 4, j % 16, k % 4], "int32"
|
|
)
|
|
|
|
|
|
@tvm.script.ir_module
|
|
class Conv2dNCHWcTIRModule:
|
|
@T.prim_func(s_tir=True)
|
|
def main(
|
|
placeholder: T.Buffer((1, 4, 56, 56, 16), "uint8"),
|
|
placeholder_1: T.Buffer((16, 4, 1, 1, 4, 16, 4), "int8"),
|
|
conv2d_NCHWc_int8: T.Buffer((1, 16, 56, 56, 16), "int32"),
|
|
) -> None:
|
|
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
|
for i0, i1, i2, i3, i4, i5, i6, i7, i8, i9 in T.grid(1, 16, 56, 56, 16, 1, 1, 4, 4, 4):
|
|
with T.sblock("conv2d_NCHWc_int8"):
|
|
(
|
|
n,
|
|
oc_chunk,
|
|
oh,
|
|
ow,
|
|
oc_block,
|
|
kh,
|
|
kw,
|
|
ic_outer,
|
|
ic_f_inner,
|
|
ic_s_inner,
|
|
) = T.axis.remap("SSSSSRRRRR", [i0, i1, i2, i3, i4, i5, i6, i7, i8, i9])
|
|
T.reads(
|
|
placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
|
|
placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
|
|
)
|
|
T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block])
|
|
with T.init():
|
|
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = 0
|
|
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = conv2d_NCHWc_int8[
|
|
n, oc_chunk, oh, ow, oc_block
|
|
] + T.cast(
|
|
placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
|
|
"int32",
|
|
) * T.cast(
|
|
placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
|
|
"int32",
|
|
)
|
|
|
|
|
|
def collect_loops(prim_func):
|
|
loops = []
|
|
|
|
def callback(node):
|
|
if isinstance(node, tvm.tirx.For):
|
|
loops.append(node)
|
|
return True
|
|
|
|
pre_order_visit(prim_func.body, callback)
|
|
|
|
return loops
|
|
|
|
|
|
def test_get_tensorize_loop_mapping_dense_16x4():
|
|
s = Schedule(DenseTIRModule)
|
|
block = s.get_sblock("compute")
|
|
|
|
info = get_tensorize_loop_mapping(s, block, dot_product_16x4_u8i8i32_desc)
|
|
|
|
assert isinstance(info, TensorizeInfo)
|
|
|
|
desc_loop_to_sref = dict((v, k) for k, v in info.loop_map.items())
|
|
|
|
desc_loops = collect_loops(dot_product_16x4_u8i8i32_desc)
|
|
_, loop_j, loop_k = s.get_loops(block)
|
|
|
|
assert desc_loops[0] in desc_loop_to_sref and desc_loops[1] in desc_loop_to_sref
|
|
assert s.get(desc_loop_to_sref[desc_loops[0]]) == s.get(loop_j)
|
|
assert s.get(desc_loop_to_sref[desc_loops[1]]) == s.get(loop_k)
|
|
|
|
|
|
def test_get_tensorize_loop_mapping_conv2d_nchwc_16x4():
|
|
s = Schedule(Conv2dNCHWcTIRModule)
|
|
block = s.get_sblock("conv2d_NCHWc_int8")
|
|
|
|
info = get_tensorize_loop_mapping(s, block, dot_product_16x4_u8i8i32_desc)
|
|
|
|
desc_loop_to_sref = dict((v, k) for k, v in info.loop_map.items())
|
|
|
|
desc_loops = collect_loops(dot_product_16x4_u8i8i32_desc)
|
|
|
|
# i4 corresonds to the inner output channel axis of the NCHWc output tensor
|
|
# for i0, i1, i2, i3, i4, i5, i6, i7, i8, i9 in T.grid(1, 16, 56, 56, 16, 1, 1, 4, 4, 4):
|
|
_, _, _, _, i4, _, _, _, _, i9 = s.get_loops(block)
|
|
|
|
assert desc_loops[0] in desc_loop_to_sref and desc_loops[1] in desc_loop_to_sref
|
|
assert s.get(desc_loop_to_sref[desc_loops[0]]) == s.get(i4)
|
|
assert s.get(desc_loop_to_sref[desc_loops[1]]) == s.get(i9)
|
|
|
|
|
|
def test_get_tensorize_loop_mapping_matmul_mma():
|
|
@T.prim_func(s_tir=True)
|
|
def matmul_16x16x16xf16f16f16_desc(
|
|
A: T.Buffer((16, 16), "float16", align=64, offset_factor=1),
|
|
B: T.Buffer((16, 16), "float16", align=64, offset_factor=1),
|
|
C: T.Buffer((16, 16), "float16", align=64, offset_factor=1),
|
|
) -> None:
|
|
with T.sblock("root"):
|
|
T.reads(C[0:16, 0:16], A[0:16, 0:16], B[0:16, 0:16])
|
|
T.writes(C[0:16, 0:16])
|
|
for i, j, k in T.grid(16, 16, 16):
|
|
with T.sblock("update"):
|
|
vii, vjj, vkk = T.axis.remap("SSR", [i, j, k])
|
|
C[vii, vjj] = C[vii, vjj] + A[vii, vkk] * B[vjj, vkk]
|
|
|
|
matmul = create_prim_func(
|
|
te_workload.matmul_relu(
|
|
n=512,
|
|
m=512,
|
|
k=512,
|
|
)
|
|
)
|
|
|
|
s = Schedule(matmul)
|
|
block = s.get_sblock("C")
|
|
i0, i1, i2 = s.get_loops(block)
|
|
desc_loops = collect_loops(matmul_16x16x16xf16f16f16_desc)
|
|
|
|
for do_reorder in [False, True]:
|
|
# Mapping should be invariant to the loop permutation
|
|
if do_reorder:
|
|
s.reorder(i2, i0, i1)
|
|
|
|
info = get_tensorize_loop_mapping(s, block, matmul_16x16x16xf16f16f16_desc)
|
|
assert info is not None
|
|
desc_loop_to_sref = dict((v, k) for k, v in info.loop_map.items())
|
|
|
|
for i in range(3):
|
|
assert desc_loops[i] in desc_loop_to_sref
|
|
|
|
assert s.get(desc_loop_to_sref[desc_loops[0]]) == s.get(i0)
|
|
assert s.get(desc_loop_to_sref[desc_loops[1]]) == s.get(i1)
|
|
assert s.get(desc_loop_to_sref[desc_loops[2]]) == s.get(i2)
|
|
|
|
|
|
def test_get_tensorize_loop_mapping_padding_matmul():
|
|
matmul = create_prim_func(
|
|
te_workload.matmul_relu(
|
|
n=127,
|
|
m=256,
|
|
k=65,
|
|
in_dtype="float16",
|
|
out_dtype="float16",
|
|
)
|
|
)
|
|
s = Schedule(matmul)
|
|
block = s.get_sblock("C")
|
|
|
|
desc = TensorIntrin.get(WMMA_SYNC_16x16x16_f16f16f16_INTRIN).desc
|
|
info = get_tensorize_loop_mapping(s, block, desc, allow_padding=True)
|
|
assert info is not None
|
|
expected_padding = [16, 1, 16]
|
|
actual_padding = info.block_iter_paddings
|
|
assert actual_padding is not None
|
|
assert len(actual_padding) == len(expected_padding)
|
|
for actual, expected in zip(actual_padding, expected_padding):
|
|
assert actual == expected
|
|
|
|
|
|
def check_index_map(workload, block_name, intrin_name, expected_index_map):
|
|
s = Schedule(workload)
|
|
block = s.get_sblock(block_name)
|
|
desc_func = TensorIntrin.get(intrin_name).desc
|
|
info = get_auto_tensorize_mapping_info(s, block, desc_func)
|
|
if expected_index_map is None:
|
|
assert info is None
|
|
return
|
|
assert len(info.mappings) == 1
|
|
assert IndexMap.from_func(expected_index_map).is_equivalent_to(info.mappings[0])
|
|
|
|
|
|
def test_get_auto_tensorize_mapping_info_conv2d():
|
|
conv2d = create_prim_func(
|
|
te_workload.conv2d_nhwc(4, 16, 16, 64, 64, 3, 1, 1, in_dtype="float16", out_dtype="float32")
|
|
)
|
|
check_index_map(
|
|
conv2d,
|
|
"conv2d_nhwc",
|
|
WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
|
|
lambda n, h, w, c, rh, rw, rc: (n * 256 + h * 16 + w, c, rh * 192 + rw * 64 + rc),
|
|
)
|
|
|
|
|
|
def test_get_auto_tensorize_mapping_info_conv2d_unit_batch():
|
|
conv2d = create_prim_func(
|
|
te_workload.conv2d_nhwc(1, 16, 16, 64, 64, 3, 1, 1, in_dtype="float16", out_dtype="float32")
|
|
)
|
|
check_index_map(
|
|
conv2d,
|
|
"conv2d_nhwc",
|
|
WMMA_SYNC_16x16x16_f16f16f32_INTRIN,
|
|
lambda n, h, w, c, rh, rw, rc: (n * 256 + h * 16 + w, c, rh * 192 + rw * 64 + rc),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("b,m,n,k", [(1, 512, 512, 512), (16, 32, 32, 32)])
|
|
def test_get_auto_tensorize_mapping_info_batch_matmul(b, m, n, k):
|
|
matmul = create_prim_func(
|
|
te_workload.batch_matmul_nkkm(b, m, n, k, in_dtype="float16", out_dtype="float32")
|
|
)
|
|
check_index_map(
|
|
matmul, "Z", WMMA_SYNC_16x16x16_f16f16f32_INTRIN, lambda b, m, n, k: (b, m, n, k)
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"n,m,k,expected",
|
|
[
|
|
(
|
|
512,
|
|
512,
|
|
512,
|
|
lambda n, m, k: (
|
|
n,
|
|
m,
|
|
k,
|
|
),
|
|
),
|
|
(1, 32, 32, lambda n, m, k: (n, m, k)),
|
|
],
|
|
)
|
|
def test_get_auto_tensorize_mapping_info_matmul(n, m, k, expected):
|
|
matmul = create_prim_func(te_workload.matmul(n, m, k, in_dtype="float16", out_dtype="float32"))
|
|
check_index_map(matmul, "C", WMMA_SYNC_16x16x16_f16f16f32_INTRIN, expected)
|
|
|
|
|
|
def test_is_output_block():
|
|
@T.prim_func(s_tir=True)
|
|
def two_elementwise(a: T.handle, c: T.handle) -> None:
|
|
A = T.match_buffer(a, (128, 128), "float32")
|
|
B = T.sblock_alloc_buffer((128, 128), "float32")
|
|
C = T.match_buffer(c, (128, 128), "float32")
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("B"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj] * 2.0
|
|
for i, j in T.grid(128, 128):
|
|
with T.sblock("C"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
C[vi, vj] = B[vi, vj] + 1.0
|
|
|
|
sch = tvm.s_tir.Schedule(two_elementwise)
|
|
block_rv = sch.get_sblock("C")
|
|
assert is_output_block(sch, block_rv)
|
|
|
|
|
|
def test_empty_grid():
|
|
@T.prim_func(s_tir=True)
|
|
def foo(out: T.Buffer((T.int64(1), T.int64(8), T.int64(8)), "int32")):
|
|
act = T.sblock_alloc_buffer((1, 8, 8), "int32")
|
|
for z2, y2, x2 in T.grid(1, 8, 8):
|
|
with T.sblock("b0"):
|
|
az, ay, ax = T.axis.remap("SSS", [z2, y2, x2])
|
|
T.writes(act[az, ay, ax])
|
|
act[az, ay, az] = T.int32(0)
|
|
# Empty grid:
|
|
for z1, y1, x1 in T.grid(0, 8, 8):
|
|
with T.sblock("b1"):
|
|
az, ay, ax = T.axis.remap("SSS", [z1, y1, x1])
|
|
T.reads(act[az + 1, ay, ax])
|
|
T.writes(out[az, ay, ax])
|
|
out[az, ay, ax] = act[az + 1, ay, ax]
|
|
# The block below is not needed to show the bug, but the 'out'
|
|
# buffer would be undefined without it.
|
|
for z2, y2, x2 in T.grid(1, 8, 8):
|
|
with T.sblock("b2"):
|
|
az, ay, ax = T.axis.remap("SSS", [z2, y2, x2])
|
|
T.writes(out[az, ay, ax])
|
|
out[az, ay, az] = T.int32(0)
|
|
|
|
# This caused a crash before.
|
|
sch = tvm.s_tir.Schedule(foo)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|