wehub-resource-sync
107 lines
3.9 KiB
Python
107 lines
3.9 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: E501, F401, F841
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
|
|
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
|
|
|
|
import tvm.topi.testing
|
|
from tvm import relax
|
|
from tvm.relax.transform import LegalizeOps
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
# TODO(tvm-team): `tirx.transform.DefaultGPUSchedule` does not work.
|
|
target, dev = "llvm", tvm.cpu()
|
|
|
|
|
|
def build(mod):
|
|
exe = tvm.compile(mod, target=target)
|
|
return relax.VirtualMachine(exe, dev)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"begin, end, strides",
|
|
[
|
|
([0, 2, 4, 4], [5, 5, 7, 8], [1, 1, 2, 3]),
|
|
([0, 2, 4, 4], [5, 5, 11, 10], [1, 1, 1, 1]),
|
|
([0, 2, 10, 14], [0, 5, 1, 1], [1, 1, -1, -2]),
|
|
],
|
|
)
|
|
def test_dynamic_strided_slice(begin, end, strides):
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class DynamicStridedSlice:
|
|
@R.function
|
|
def main(x: R.Tensor((8, 9, 10, 10), "float32"), begin: R.Tensor((4,),"int64"), end: R.Tensor((4,),"int64"), strides: R.Tensor((4,),"int64")) -> R.Tensor("float32", ndim=4):
|
|
gv: R.Tensor("float32", ndim=4) = R.dynamic_strided_slice(x, begin, end, strides)
|
|
return gv
|
|
# fmt: on
|
|
vm = build(DynamicStridedSlice)
|
|
|
|
x_np = np.random.rand(8, 9, 10, 10).astype(np.float32)
|
|
data_nd = tvm.runtime.tensor(x_np, dev)
|
|
begin_nd = tvm.runtime.tensor(np.array(begin).astype("int64"), dev)
|
|
end_nd = tvm.runtime.tensor(np.array(end).astype("int64"), dev)
|
|
strides_nd = tvm.runtime.tensor(np.array(strides).astype("int64"), dev)
|
|
|
|
# Reference implementation
|
|
out_npy = tvm.topi.testing.strided_slice_python(x_np, begin, end, strides)
|
|
out_nd = vm["main"](data_nd, begin_nd, end_nd, strides_nd)
|
|
tvm.testing.assert_allclose(out_nd.numpy(), out_npy)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"begin, end, strides",
|
|
[
|
|
([0, 2, 4, 4], [5, 5, 7, 8], [1, 1, 2, 3]),
|
|
([0, 2, 4, 4], [5, 5, 11, 10], [1, 1, 1, 1]),
|
|
([0, 2, 10, 14], [0, 5, 1, 1], [1, 1, -1, -2]),
|
|
],
|
|
)
|
|
def test_dynamic_strided_slice_symbolic(begin, end, strides):
|
|
# fmt: off
|
|
@tvm.script.ir_module
|
|
class DynamicStridedSlice:
|
|
@R.function
|
|
def main(x: R.Tensor(("m", "n", 10, 10), "float32"), begin: R.Tensor((4,),"int64"), end: R.Tensor((4,),"int64"), strides: R.Tensor((4,),"int64")) -> R.Tensor("float32", ndim=4):
|
|
m = T.int64()
|
|
n = T.int64()
|
|
gv: R.Tensor("float32", ndim=4) = R.dynamic_strided_slice(x, begin, end, strides)
|
|
return gv
|
|
# fmt: on
|
|
vm = build(DynamicStridedSlice)
|
|
|
|
x_np = np.random.rand(8, 9, 10, 10).astype(np.float32)
|
|
data_nd = tvm.runtime.tensor(x_np, dev)
|
|
begin_nd = tvm.runtime.tensor(np.array(begin).astype("int64"), dev)
|
|
end_nd = tvm.runtime.tensor(np.array(end).astype("int64"), dev)
|
|
strides_nd = tvm.runtime.tensor(np.array(strides).astype("int64"), dev)
|
|
|
|
# Reference implementation
|
|
out_npy = tvm.topi.testing.strided_slice_python(x_np, begin, end, strides)
|
|
out_nd = vm["main"](data_nd, begin_nd, end_nd, strides_nd)
|
|
tvm.testing.assert_allclose(out_nd.numpy(), out_npy)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|