| # 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. |
| |
| from typing import Optional, Union |
| |
| import tvm |
| import tvm.script |
| import tvm.testing |
| from tvm import IRModule, relax |
| from tvm.script import relax as R, ir as I |
| |
| |
| def _check( |
| parsed: Union[relax.Function, IRModule], |
| expect: Optional[Union[relax.Function, IRModule]], |
| ): |
| test = parsed.script(show_meta=True) |
| roundtrip_mod = tvm.script.from_source(test) |
| tvm.ir.assert_structural_equal(parsed, roundtrip_mod) |
| if expect: |
| tvm.ir.assert_structural_equal(parsed, expect) |
| |
| |
| def test_broadcast_to(): |
| @R.function |
| def foo(x: R.Tensor((2, 1, 3), "float32")) -> R.Tensor((4, 2, 5, 3), "float32"): |
| gv: R.Tensor((4, 2, 5, 3), "float32") = R.broadcast_to(x, (4, 2, 5, 3)) |
| return gv |
| |
| bb = relax.BlockBuilder() |
| x = relax.Var("x", R.Tensor((2, 1, 3), "float32")) |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.broadcast_to(x, (4, 2, 5, 3))) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_concat(): |
| @R.function |
| def foo( |
| x1: R.Tensor((1, 2, 3), "float32"), |
| x2: R.Tensor((1, 3, 3), "float32"), |
| x3: R.Tensor((1, 4, 3), "float32"), |
| ) -> R.Tensor((1, 9, 3), "float32"): |
| gv: R.Tensor((1, 9, 3), "float32") = R.concat((x1, x2, x3), axis=1) |
| return gv |
| |
| x1 = relax.Var("x1", R.Tensor((1, 2, 3), "float32")) |
| x2 = relax.Var("x2", R.Tensor((1, 3, 3), "float32")) |
| x3 = relax.Var("x3", R.Tensor((1, 4, 3), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x1, x2, x3]): |
| gv = bb.emit(relax.op.concat((x1, x2, x3), axis=1)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_concat_without_specified_axis(): |
| @R.function |
| def foo( |
| x1: R.Tensor((2,), "float32"), x2: R.Tensor((3,), "float32"), x3: R.Tensor((4,), "float32") |
| ) -> R.Tensor((9,), "float32"): |
| gv: R.Tensor((9,), "float32") = R.concat((x1, x2, x3), axis=None) |
| return gv |
| |
| x1 = relax.Var("x1", R.Tensor((2,), "float32")) |
| x2 = relax.Var("x2", R.Tensor((3,), "float32")) |
| x3 = relax.Var("x3", R.Tensor((4,), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x1, x2, x3]): |
| gv = bb.emit(relax.op.concat((x1, x2, x3), axis=None)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_expand_dims(): |
| @R.function |
| def foo(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), "float32"): |
| gv: R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), "float32") = R.expand_dims(x, axis=[-1, 1, -6, 3, 5]) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((2, 3, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.expand_dims(x, axis=[-1, 1, -6, 3, 5])) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_flatten(): |
| @R.function |
| def foo(x: R.Tensor((3, 4, 5), "float32")) -> R.Tensor((60,), "float32"): |
| gv: R.Tensor((60,), "float32") = R.flatten(x) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((3, 4, 5), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.flatten(x)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_layout_transform(): |
| transformation = lambda n, c, h, w: (n, h, w, c) |
| |
| @R.function |
| def foo(x: R.Tensor((2, 3, 4, 5), "float32")): |
| gv: R.Tensor((2, 4, 5, 3), "float32") = R.layout_transform(x, index_map=transformation) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((2, 3, 4, 5), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.layout_transform(x, index_map=transformation)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_layout_transform_with_padding(): |
| transformation = lambda n, c, h, w: (n, c // 3, h, w, c % 3) |
| |
| @R.function |
| def foo(x: R.Tensor((10, 20, 2, 2), "float32")): |
| gv: R.Tensor((10, 7, 2, 2, 3), "float32") = R.layout_transform( |
| x, index_map=transformation, pad_value=2 |
| ) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((10, 20, 2, 2), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.layout_transform(x, index_map=transformation, pad_value=2)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_permute_dims(): |
| @R.function |
| def foo(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((2, 4, 3, 1), "float32"): |
| gv: R.Tensor((2, 4, 3, 1), "float32") = R.permute_dims(x, axes=[1, -1, 2, -4]) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((1, 2, 3, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.permute_dims(x, axes=[1, -1, 2, -4])) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_permute_dims_none_arg(): |
| @R.function |
| def foo(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((4, 3, 2, 1), "float32"): |
| gv: R.Tensor((4, 3, 2, 1), "float32") = R.permute_dims(x) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((1, 2, 3, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.permute_dims(x)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_reshape(): |
| @R.function |
| def foo(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((8, 3), "float32"): |
| gv: R.Tensor((8, 3), "float32") = R.reshape(x, (8, 3)) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((1, 2, 3, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.reshape(x, shape=(8, 3))) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_reshape_infer_dim(): |
| @R.function |
| def foo(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((8, 1, 3), "float32"): |
| gv: R.Tensor((8, 1, 3), "float32") = R.reshape(x, (8, -1, 3)) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((1, 2, 3, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.reshape(x, shape=(8, -1, 3))) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_split_by_indices(): |
| @R.function |
| def foo( |
| x: R.Tensor((2, 10, 4), dtype="float32") |
| ) -> R.Tuple( |
| R.Tensor((2, 0, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 4, 4), dtype="float32"), |
| R.Tensor((2, 0, 4), dtype="float32"), |
| R.Tensor((2, 4, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 0, 4), dtype="float32"), |
| R.Tensor((2, 1, 4), dtype="float32"), |
| ): |
| gv: R.Tuple( |
| R.Tensor((2, 0, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 4, 4), dtype="float32"), |
| R.Tensor((2, 0, 4), dtype="float32"), |
| R.Tensor((2, 4, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 0, 4), dtype="float32"), |
| R.Tensor((2, 1, 4), dtype="float32"), |
| ) = R.split(x, indices_or_sections=[-2, 2, 6, 4, 8, 12, 9], axis=1) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((2, 10, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.split(x, indices_or_sections=[-2, 2, 6, 4, 8, 12, 9], axis=1)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_split_by_n_section(): |
| @R.function |
| def foo( |
| x: R.Tensor((2, 10, 4), dtype="float32") |
| ) -> R.Tuple( |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| ): |
| gv: R.Tuple( |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| R.Tensor((2, 2, 4), dtype="float32"), |
| ) = R.split(x, indices_or_sections=5, axis=1) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((2, 10, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.split(x, indices_or_sections=5, axis=1)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_squeeze(): |
| @R.function |
| def foo(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) -> R.Tensor((2, 3, 4), "float32"): |
| gv: R.Tensor((2, 3, 4), "float32") = R.squeeze(x) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((2, 1, 3, 1, 1, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.squeeze(x)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_squeeze_with_indices(): |
| @R.function |
| def foo(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) -> R.Tensor((2, 3, 1, 4), "float32"): |
| gv: R.Tensor((2, 3, 1, 4), "float32") = R.squeeze(x, axis=[3, -5]) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((2, 1, 3, 1, 1, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.squeeze(x, axis=[3, -5])) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_collapse_sum_like(): |
| @R.function |
| def foo( |
| x: R.Tensor((3, 4, 5), "float32"), y: R.Tensor((4, 5), "float32") |
| ) -> R.Tensor((4, 5), "float32"): |
| gv: R.Tensor((4, 5), "float32") = R.collapse_sum_like(x, y) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((3, 4, 5), "float32")) |
| y = relax.Var("y", R.Tensor((4, 5), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x, y]): |
| gv = bb.emit(relax.op.collapse_sum_like(x, y)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_collapse_sum_to(): |
| @R.function |
| def foo(x: R.Tensor((3, 4, 5), "float32")) -> R.Tensor((4, 5), "float32"): |
| gv: R.Tensor((4, 5), "float32") = R.collapse_sum_to(x, (4, 5)) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((3, 4, 5), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.collapse_sum_to(x, (4, 5))) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_repeat(): |
| @R.function |
| def foo(x: R.Tensor((2, 3, 4), "float32")): |
| gv = R.repeat(x, 3, 1) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((2, 3, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.repeat(x, 3, 1)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_repeat_no_axis(): |
| @R.function |
| def foo(x: R.Tensor((2, 3, 4), "float32")): |
| gv = R.repeat(x, 3) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((2, 3, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.repeat(x, 3)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_tile(): |
| @R.function |
| def foo(x: R.Tensor((2, 3, 4), "float32")): |
| gv = R.tile(x, (2, 3)) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((2, 3, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.tile(x, (2, 3))) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_flip(): |
| @R.function |
| def foo(x: R.Tensor((2, 3, 4), "float32")): |
| gv = R.flip(x, axis=1) |
| return gv |
| |
| x = relax.Var("x", R.Tensor((2, 3, 4), "float32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", [x]): |
| gv = bb.emit(relax.op.flip(x, axis=1)) |
| bb.emit_func_output(gv) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| def test_to_vdevice(): |
| @I.ir_module |
| class ToVDevice: |
| I.module_global_infos({"vdevice": [I.vdevice("llvm")]}) |
| |
| @R.function |
| def foo(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"): |
| tensor = R.to_vdevice(x, "llvm") |
| return tensor |
| |
| x = relax.Var("x", R.Tensor((), "int32")) |
| bb = relax.BlockBuilder() |
| vdev = I.vdevice("llvm") |
| with bb.function("foo", (x,)): |
| tensor = bb.emit(relax.op.to_vdevice(x, vdev)) |
| bb.emit_func_output(tensor) |
| bb.get().update_global_info("vdevice", [vdev]) |
| |
| _check(ToVDevice, bb.get()) |
| |
| |
| def test_hint_on_device(): |
| @R.function |
| def foo(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"): |
| r = R.hint_on_device(x, R.device(1, 0)) |
| return r |
| |
| x = relax.Var("x", R.Tensor((), "int32")) |
| bb = relax.BlockBuilder() |
| with bb.function("foo", (x,)): |
| tensor = bb.emit(relax.op.hint_on_device(x, R.cpu())) |
| bb.emit_func_output(tensor) |
| |
| _check(foo, bb.get()["foo"]) |
| |
| |
| if __name__ == "__main__": |
| tvm.testing.main() |
