wehub-resource-sync
369 lines
14 KiB
Python
369 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.
|
|
# ruff: noqa: E741, F401
|
|
import numpy as np
|
|
import pytest
|
|
import tvm_ffi
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.ir import assert_structural_equal
|
|
from tvm.runtime import const
|
|
from tvm.script import tirx as T
|
|
from tvm.tirx import IndexMap, IntImm, floordiv, floormod, stmt_functor
|
|
|
|
|
|
def assert_equal_index_map(map1: IndexMap, map2: IndexMap) -> None:
|
|
iters_1 = map1.map_indices(map2.initial_indices)
|
|
iters_2 = map2.final_indices
|
|
assert len(iters_1) == len(iters_2)
|
|
|
|
analyzer = tvm.arith.Analyzer()
|
|
for iter1, iter2 in zip(iters_1, iters_2):
|
|
assert analyzer.can_prove_equal(iter1, iter2)
|
|
|
|
|
|
def test_index_mapping():
|
|
index_map = IndexMap.from_func(lambda i: [i // 4, i % 4], index_dtype="int32")
|
|
|
|
assert_structural_equal(index_map.map_indices([0]), [T.int32(0), T.int32(0)])
|
|
assert_structural_equal(index_map.map_indices([3]), [T.int32(0), T.int32(3)])
|
|
assert_structural_equal(index_map.map_indices([4]), [T.int32(1), T.int32(0)])
|
|
assert_structural_equal(index_map.map_indices([42]), [T.int32(10), T.int32(2)])
|
|
assert_structural_equal(index_map.map_indices([T.int64(42)]), [T.int64(10), T.int64(2)])
|
|
|
|
|
|
def test_map_indices_accepts_external_analyzer():
|
|
tile = tvm.tirx.Var("tile", "int32")
|
|
index_map = IndexMap.from_func(lambda i: [i // tile], index_dtype="int32")
|
|
analyzer = tvm.arith.Analyzer()
|
|
|
|
unsimplified = index_map.map_indices([T.int32(32)])[0]
|
|
analyzer.bind(tile, T.int32(16))
|
|
simplified = index_map.map_indices([T.int32(32)], analyzer=analyzer)[0]
|
|
|
|
assert not tvm_ffi.structural_equal(unsimplified, T.int32(2))
|
|
assert_structural_equal(simplified, T.int32(2))
|
|
|
|
|
|
def test_map_shape_accepts_external_analyzer():
|
|
tile = tvm.tirx.Var("tile", "int32")
|
|
index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32")
|
|
analyzer = tvm.arith.Analyzer()
|
|
|
|
analyzer.bind(tile, T.int32(16))
|
|
mapped_shape = index_map.map_shape([T.int32(32)], analyzer=analyzer)
|
|
|
|
assert_structural_equal(mapped_shape, [T.int32(2), T.int32(16)])
|
|
|
|
|
|
def test_is_equivalent_to_accepts_external_analyzer():
|
|
tile = tvm.tirx.Var("tile", "int32")
|
|
concrete = IndexMap.from_func(lambda i: [i // 4, i % 4], index_dtype="int32")
|
|
symbolic = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32")
|
|
|
|
# Without binding `tile`, the symbolic map cannot be proven equivalent.
|
|
assert not concrete.is_equivalent_to(symbolic)
|
|
|
|
analyzer = tvm.arith.Analyzer()
|
|
analyzer.bind(tile, T.int32(4))
|
|
assert concrete.is_equivalent_to(symbolic, analyzer=analyzer)
|
|
|
|
|
|
def test_shape_mapping():
|
|
index_map = IndexMap.from_func(lambda i: [i // 4, i % 4], index_dtype="int32")
|
|
|
|
assert_structural_equal(index_map.map_shape([4]), [T.int32(1), T.int32(4)])
|
|
assert_structural_equal(index_map.map_shape([16]), [T.int32(4), T.int32(4)])
|
|
|
|
assert_structural_equal(index_map.map_shape([14]), [T.int32(4), T.int32(4)])
|
|
assert_structural_equal(index_map.map_shape([T.int64(16)]), [T.int64(4), T.int64(4)])
|
|
assert_structural_equal(index_map.map_shape([T.int64(14)]), [T.int64(4), T.int64(4)])
|
|
|
|
|
|
def test_inverse():
|
|
index_map = IndexMap.from_func(lambda i: [i // 4, i % 4])
|
|
expected_inverse = IndexMap.from_func(lambda i, j: [4 * i + j])
|
|
|
|
assert index_map.inverse([16]).is_equivalent_to(expected_inverse)
|
|
|
|
|
|
def test_inverse_preserves_passthrough_var_names():
|
|
index_map = IndexMap.from_func(lambda i, j: [j, i], index_dtype="int32")
|
|
inverse = index_map.inverse([8, 16])
|
|
assert [v.name for v in inverse.initial_indices] == ["j", "i"]
|
|
|
|
|
|
def test_inverse_accepts_external_analyzer():
|
|
tile = tvm.tirx.Var("tile", "int32")
|
|
index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32")
|
|
analyzer = tvm.arith.Analyzer()
|
|
|
|
analyzer.bind(tile, T.int32(16))
|
|
inverse = index_map.inverse([T.int32(32)], analyzer=analyzer)
|
|
mapped = inverse.map_indices([T.int32(1), T.int32(3)], analyzer=analyzer)
|
|
|
|
assert_structural_equal(mapped, [T.int32(19)])
|
|
|
|
|
|
def test_nonbijective_inverse_gives_error():
|
|
index_map = IndexMap.from_func(lambda i: [i // 4, i % 4])
|
|
|
|
with pytest.raises(RuntimeError):
|
|
index_map.inverse([14])
|
|
|
|
|
|
dynamic_N = tvm.tirx.Var("N", "int32")
|
|
padding_test_case = tvm.testing.parameter(
|
|
by_dict={
|
|
"no_padding": dict(
|
|
forward=lambda i: [i // 4, i % 4],
|
|
inverse=lambda i, j: [4 * i + j],
|
|
pre_shape=[16],
|
|
post_shape=[T.int32(4), T.int32(4)],
|
|
padding=lambda i, j: tvm.runtime.convert(False),
|
|
),
|
|
"right_padding": dict(
|
|
forward=lambda i: [i // 4, i % 4],
|
|
inverse=lambda i, j: [4 * i + j],
|
|
pre_shape=[15],
|
|
post_shape=[T.int32(4), T.int32(4)],
|
|
padding=lambda i, j: tvm.tirx.And(i == 3, tvm.runtime.convert(3) == j),
|
|
),
|
|
"left_padding": dict(
|
|
forward=lambda i: [(i + 1) // 4, (i + 1) % 4],
|
|
inverse=lambda i, j: [4 * i + j - 1],
|
|
pre_shape=[15],
|
|
post_shape=[T.int32(4), T.int32(4)],
|
|
padding=lambda i, j: tvm.tirx.And(i == 0, j < 1),
|
|
),
|
|
"left_and_right_padding": dict(
|
|
forward=lambda i: [(i + 1) // 4, (i + 1) % 4],
|
|
inverse=lambda i, j: [4 * i + j - 1],
|
|
pre_shape=[14],
|
|
post_shape=[T.int32(4), T.int32(4)],
|
|
padding=lambda i, j: tvm.tirx.Or(
|
|
tvm.tirx.And(i == 0, j < 1),
|
|
tvm.tirx.And(i == 3, tvm.runtime.convert(3) == j),
|
|
),
|
|
),
|
|
"dynamic_size": dict(
|
|
forward=lambda i: [i // 4, i % 4],
|
|
inverse=lambda i, j: [4 * i + j],
|
|
pre_shape=[dynamic_N],
|
|
post_shape=[(dynamic_N - dynamic_N % (-4)) // 4, T.int32(4)],
|
|
padding=lambda i, j: tvm.tirx.And(
|
|
dynamic_N % (-4) != 0,
|
|
tvm.tirx.And(i == dynamic_N // 4, j >= dynamic_N % 4),
|
|
),
|
|
),
|
|
"2d_padding": dict(
|
|
forward=lambda i, j: [(i + 1) // 4, (j + 5) // 8, (i + 1) % 4, (j + 5) % 8],
|
|
inverse=lambda i_outer, j_outer, i_inner, j_inner: [
|
|
4 * i_outer + i_inner - 1,
|
|
8 * j_outer + j_inner - 5,
|
|
],
|
|
pre_shape=[14, 31],
|
|
post_shape=[
|
|
T.int32(4), # ceildiv(left_pad + i.extent, 4) = ceildiv(1 + 14, 4) = 4
|
|
T.int32(5), # ceildiv(left_pad + j.extent, 8) = ceildiv(5 + 31, 8) = 5
|
|
T.int32(4), # Range of iter%4
|
|
T.int32(8), # Range of iter%8
|
|
],
|
|
padding=lambda i_outer, j_outer, i_inner, j_inner: tvm.tirx.Or(
|
|
tvm.tirx.Or(
|
|
tvm.tirx.And(i_outer == 0, i_inner < 1),
|
|
tvm.tirx.And(i_outer == 3, tvm.runtime.convert(3) == i_inner),
|
|
),
|
|
tvm.tirx.Or(
|
|
tvm.tirx.And(j_outer == 0, j_inner < 5),
|
|
tvm.tirx.And(j_outer == 4, j_inner >= 4),
|
|
),
|
|
),
|
|
),
|
|
"multiple_right_padding": dict(
|
|
forward=lambda i: [i // 32, (i // 4) % 8, i % 4],
|
|
inverse=lambda i, j, k: [32 * i + 4 * j + k],
|
|
pre_shape=[116],
|
|
post_shape=[T.int32(4), T.int32(8), T.int32(4)],
|
|
padding=lambda i, j, k: tvm.tirx.And(i == 3, 4 * j + k >= 20),
|
|
),
|
|
"multiple_right_padding_transpose": dict(
|
|
forward=lambda i: [(i // 4) % 8, i // 32, i % 4],
|
|
inverse=lambda j, i, k: [32 * i + 4 * j + k],
|
|
pre_shape=[116],
|
|
post_shape=[T.int32(8), T.int32(4), T.int32(4)],
|
|
padding=lambda j, i, k: tvm.tirx.And(i == 3, 4 * j + k >= 20),
|
|
),
|
|
"multiple_left_padding": dict(
|
|
forward=lambda i: [(i + 5) // 32, ((i + 5) // 4) % 8, (i + 5) % 4],
|
|
inverse=lambda i, j, k: [32 * i + 4 * j + k - 5],
|
|
pre_shape=[123],
|
|
post_shape=[T.int32(4), T.int32(8), T.int32(4)],
|
|
padding=lambda i, j, k: tvm.tirx.And(i == 0, j * 4 + k < 5),
|
|
),
|
|
"multiple_left_padding_with_transpose": dict(
|
|
forward=lambda i: [((i + 5) // 4) % 8, (i + 5) // 32, (i + 5) % 4],
|
|
inverse=lambda j, i, k: [32 * i + 4 * j + k - 5],
|
|
pre_shape=[123],
|
|
post_shape=[T.int32(8), T.int32(4), T.int32(4)],
|
|
padding=lambda j, i, k: tvm.tirx.And(i == 0, j * 4 + k < 5),
|
|
),
|
|
"outer_loop_extent_one": dict(
|
|
forward=lambda i: [i // 4, i % 4],
|
|
inverse=lambda i, j: [i * 4 + j],
|
|
pre_shape=[3],
|
|
post_shape=[T.int32(1), T.int32(4)],
|
|
padding=lambda i, j: tvm.runtime.convert(3) == j,
|
|
),
|
|
}
|
|
)
|
|
|
|
|
|
def test_nonsurjective_inverse(padding_test_case):
|
|
index_map = IndexMap.from_func(padding_test_case["forward"], index_dtype="int32")
|
|
|
|
inverse, padding_predicate = index_map.non_surjective_inverse(padding_test_case["pre_shape"])
|
|
expected_inverse = IndexMap.from_func(padding_test_case["inverse"])
|
|
assert inverse.is_equivalent_to(expected_inverse)
|
|
|
|
post_shape = index_map.map_shape(padding_test_case["pre_shape"])
|
|
tvm.ir.assert_structural_equal(post_shape, padding_test_case["post_shape"])
|
|
|
|
expected_predicate = padding_test_case["padding"](*inverse.initial_indices)
|
|
|
|
# Can't use analyzer.can_prove_equal, because it can't simplify
|
|
# expressions like `(4*i+j >= 14) - (4*i+j >= 14)`.
|
|
analyzer = tvm.arith.Analyzer()
|
|
expected_predicate = analyzer.simplify(expected_predicate)
|
|
padding_predicate = analyzer.simplify(padding_predicate)
|
|
tvm.ir.assert_structural_equal(padding_predicate, expected_predicate)
|
|
|
|
|
|
def test_non_surjective_inverse_accepts_external_analyzer():
|
|
tile = tvm.tirx.Var("tile", "int32")
|
|
index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32")
|
|
analyzer = tvm.arith.Analyzer()
|
|
|
|
analyzer.bind(tile, T.int32(16))
|
|
inverse, padding_predicate = index_map.non_surjective_inverse([T.int32(31)], analyzer=analyzer)
|
|
mapped = inverse.map_indices([T.int32(1), T.int32(15)], analyzer=analyzer)
|
|
|
|
assert_structural_equal(mapped, [T.int32(31)])
|
|
|
|
padding_at_last_element = stmt_functor.substitute(
|
|
padding_predicate,
|
|
{inverse.initial_indices[0]: T.int32(1), inverse.initial_indices[1]: T.int32(15)},
|
|
)
|
|
padding_at_first_element = stmt_functor.substitute(
|
|
padding_predicate,
|
|
{inverse.initial_indices[0]: T.int32(0), inverse.initial_indices[1]: T.int32(0)},
|
|
)
|
|
assert_structural_equal(analyzer.simplify(padding_at_last_element), T.bool(True))
|
|
assert_structural_equal(analyzer.simplify(padding_at_first_element), T.bool(False))
|
|
|
|
|
|
def test_non_surjective_inverse_does_not_bind_output_vars_to_external_analyzer():
|
|
tile = tvm.tirx.Var("tile", "int32")
|
|
index_map = IndexMap.from_func(lambda i: [i // tile, i % tile], index_dtype="int32")
|
|
analyzer = tvm.arith.Analyzer()
|
|
|
|
analyzer.bind(tile, T.int32(16))
|
|
inverse, _ = index_map.non_surjective_inverse([T.int32(31)], analyzer=analyzer)
|
|
|
|
analyzer.bind(inverse.initial_indices[0], T.int32(0))
|
|
analyzer.bind(inverse.initial_indices[1], T.int32(1))
|
|
|
|
|
|
def test_index_map_inverse_no_iter():
|
|
def input_example(i0, i1, i2, i3):
|
|
j0 = floordiv(i3, 32)
|
|
j1 = floordiv(i2, 2)
|
|
j2 = floormod(i2, 2)
|
|
j3 = floormod(i3, 32)
|
|
return j0, j1, j2, j3
|
|
|
|
def expected_inverse(i0, i1, i2, i3):
|
|
return IntImm("int32", 0), IntImm("int32", 0), i2 + i1 * 2, i3 + i0 * 32
|
|
|
|
index_map = IndexMap.from_func(input_example)
|
|
inverse_map = index_map.inverse([1, 1, 64, 64])
|
|
expected_map = IndexMap.from_func(expected_inverse)
|
|
assert expected_map.is_equivalent_to(inverse_map)
|
|
|
|
|
|
def test_map_tensor():
|
|
index_map = IndexMap.from_func(lambda i: [i // 4, i % 4])
|
|
|
|
inp = np.arange(16).astype("int8")
|
|
|
|
out = index_map.map_tensor(tvm.runtime.tensor(inp)).numpy()
|
|
|
|
ref = np.zeros(out.shape).astype("int8")
|
|
|
|
for i in range(16):
|
|
ref[i // 4, i % 4] = inp[i]
|
|
|
|
np.testing.assert_equal(ref, out)
|
|
|
|
index_map = IndexMap.from_func(lambda i0, i1, i2, i3: (i3, i0, i1, i2))
|
|
|
|
inp = np.random.randn(10, 10, 10, 10).astype("float16")
|
|
|
|
out = index_map.map_tensor(tvm.runtime.tensor(inp)).numpy()
|
|
|
|
ref = np.transpose(inp, (3, 0, 1, 2))
|
|
|
|
np.testing.assert_equal(ref, out)
|
|
|
|
index_map = IndexMap.from_func(
|
|
lambda i0, i1, i2, i3: (
|
|
floordiv(i3, 32),
|
|
i0,
|
|
floordiv(i2, 8),
|
|
floordiv(floormod(i3, 32), 16),
|
|
i1,
|
|
floormod(i2, 8),
|
|
floormod(i3, 16),
|
|
)
|
|
)
|
|
|
|
kH = kW = 3
|
|
I = 64
|
|
O = 64
|
|
inp = np.random.randn(kH, kW, I, O).astype("float32")
|
|
arr = tvm.runtime.tensor(inp)
|
|
out = index_map.map_tensor(arr).numpy()
|
|
|
|
ref = np.zeros(out.shape).astype("float32")
|
|
|
|
for i0 in range(kH):
|
|
for i1 in range(kW):
|
|
for i2 in range(I):
|
|
for i3 in range(O):
|
|
v = inp[i0, i1, i2, i3]
|
|
ref[i3 // 32, i0, i2 // 8, (i3 % 32) // 16, i1, i2 % 8, i3 % 16] = v
|
|
|
|
np.testing.assert_equal(ref, out)
|
|
|
|
inverse_map = index_map.inverse(inp.shape)
|
|
np.testing.assert_equal(inverse_map.map_tensor(index_map.map_tensor(arr)).numpy(), inp)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|