wehub-resource-sync
225 lines
7.4 KiB
Python
225 lines
7.4 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.
|
|
"""Test the TIR codegen path of VM compiled mode.
|
|
|
|
Restrictions: all shape lowered, explicit allocation.
|
|
"""
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.ir import assert_structural_equal
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def get_tir_mod(mod):
|
|
builder = relax.ExecBuilder()
|
|
return relax.vm_build._vmcodegen(builder, mod, exec_mode="compiled")
|
|
|
|
|
|
def test_add():
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function(pure=False)
|
|
def foo(x: R.Tensor):
|
|
R.func_attr({"global_symbol": "foo"})
|
|
z = R.call_packed("test.vm.add", x, x, ty_args=(R.Tensor))
|
|
return z
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def __vmtir__foo(ctx_ptr: T.handle, r: T.handle, c: T.handle, f: T.handle):
|
|
T.func_attr({"global_symbol": "__vmtir__foo"})
|
|
T.anylist_setitem_call_packed(
|
|
r,
|
|
T.int32(2),
|
|
"test.vm.add",
|
|
T.anylist_getitem(r, T.int32(0)),
|
|
T.anylist_getitem(r, T.int32(0)),
|
|
)
|
|
T.anylist_setitem_call_packed(
|
|
r, T.int32(1), "vm.builtin.copy", T.anylist_getitem(r, T.int32(2))
|
|
)
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = get_tir_mod(before)
|
|
assert_structural_equal(expected, after)
|
|
|
|
|
|
def test_tir_call():
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@T.prim_func(s_tir=True)
|
|
def shape_func(H: T.Buffer(T.int64(4), "int64")):
|
|
T.func_attr({"global_symbol": "shape_func"})
|
|
# generated compute function
|
|
H[T.int64(0)] = H[T.int64(0)] + T.int64(1)
|
|
|
|
@R.function(pure=False)
|
|
def foo(x: R.Tensor([4], "int64")):
|
|
R.func_attr({"global_symbol": "foo"})
|
|
_ = Before.shape_func(x)
|
|
return x
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def shape_func(H: T.Buffer(T.int64(4), "int64")):
|
|
T.func_attr({"global_symbol": "shape_func"})
|
|
# generated compute function
|
|
H[T.int64(0)] = H[T.int64(0)] + T.int64(1)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def __vmtir__foo(ctx_ptr: T.handle, r: T.handle, c: T.handle, f: T.handle):
|
|
T.func_attr({"global_symbol": "__vmtir__foo"})
|
|
T.call_cpacked("shape_func", T.anylist_getitem(r, T.int32(0)))
|
|
T.anylist_setitem_call_packed(
|
|
r, T.int32(1), "vm.builtin.copy", T.anylist_getitem(r, T.int32(0))
|
|
)
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = get_tir_mod(before)
|
|
assert_structural_equal(expected, after)
|
|
|
|
|
|
def test_if_cond():
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function(pure=False)
|
|
def ife(cond: R.Tensor((), "bool"), x: R.Tensor) -> R.Tensor:
|
|
R.func_attr({"global_symbol": "ife"})
|
|
if cond:
|
|
w = R.call_packed("test.vm.add", x, x, ty_args=(R.Tensor))
|
|
else:
|
|
w = R.call_packed("test.vm.mul", x, x, ty_args=(R.Tensor))
|
|
return w
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def __vmtir__ife(ctx_ptr: T.handle, r: T.handle, c: T.handle, f: T.handle):
|
|
T.func_attr({"global_symbol": "__vmtir__ife"})
|
|
if T.Call(
|
|
tvm.ir.Op.get("tirx.tvm_call_packed"),
|
|
["vm.builtin.read_if_cond", T.anylist_getitem(r, T.int32(0))],
|
|
ret_ty="bool",
|
|
):
|
|
T.anylist_setitem_call_packed(
|
|
r,
|
|
T.int32(4),
|
|
"test.vm.add",
|
|
T.anylist_getitem(r, T.int32(1)),
|
|
T.anylist_getitem(r, T.int32(1)),
|
|
)
|
|
T.anylist_setitem_call_packed(
|
|
r, T.int32(3), "vm.builtin.copy", T.anylist_getitem(r, T.int32(4))
|
|
)
|
|
else:
|
|
T.anylist_setitem_call_packed(
|
|
r,
|
|
T.int32(5),
|
|
"test.vm.mul",
|
|
T.anylist_getitem(r, T.int32(1)),
|
|
T.anylist_getitem(r, T.int32(1)),
|
|
)
|
|
T.anylist_setitem_call_packed(
|
|
r, T.int32(3), "vm.builtin.copy", T.anylist_getitem(r, T.int32(5))
|
|
)
|
|
T.anylist_setitem_call_packed(
|
|
r, T.int32(2), "vm.builtin.copy", T.anylist_getitem(r, T.int32(3))
|
|
)
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = get_tir_mod(before)
|
|
assert_structural_equal(expected, after)
|
|
|
|
|
|
def test_const():
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor):
|
|
R.func_attr({"global_symbol": "main"})
|
|
y = R.const([1, 2])
|
|
z = (y, R.const([3, 4]), x)
|
|
return z
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def __vmtir__main(ctx_ptr: T.handle, r: T.handle, c: T.handle, f: T.handle):
|
|
# function attr dict
|
|
T.func_attr({"global_symbol": "__vmtir__main"})
|
|
# body
|
|
T.anylist_setitem_call_packed(
|
|
r,
|
|
T.int32(2),
|
|
"vm.builtin.make_tuple",
|
|
T.anylist_getitem(c, T.int32(0)),
|
|
T.anylist_getitem(c, T.int32(1)),
|
|
T.anylist_getitem(r, T.int32(0)),
|
|
)
|
|
T.anylist_setitem_call_packed(
|
|
r, T.int32(1), "vm.builtin.copy", T.anylist_getitem(r, T.int32(2))
|
|
)
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = get_tir_mod(before)
|
|
assert_structural_equal(expected, after)
|
|
|
|
|
|
def test_const_call():
|
|
@tvm.script.ir_module
|
|
class Before:
|
|
@R.function(pure=False)
|
|
def main(x: R.Tensor):
|
|
R.func_attr({"global_symbol": "main"})
|
|
y = R.const([1, 2])
|
|
z = R.call_packed("test.vm.add", x, y, ty_args=(R.Tensor))
|
|
return z
|
|
|
|
@tvm.script.ir_module
|
|
class Expected:
|
|
@T.prim_func(s_tir=True)
|
|
def __vmtir__main(ctx_ptr: T.handle, r: T.handle, c: T.handle, f: T.handle):
|
|
# function attr dict
|
|
T.func_attr({"global_symbol": "__vmtir__main"})
|
|
# body
|
|
T.anylist_setitem_call_packed(
|
|
r,
|
|
2,
|
|
"test.vm.add",
|
|
T.anylist_getitem(r, 0),
|
|
T.anylist_getitem(c, 0),
|
|
)
|
|
T.anylist_setitem_call_packed(r, 1, "vm.builtin.copy", T.anylist_getitem(r, 2))
|
|
|
|
before = Before
|
|
expected = Expected
|
|
after = get_tir_mod(before)
|
|
assert_structural_equal(expected, after)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|