wehub-resource-sync
153 lines
4.9 KiB
Python
153 lines
4.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.
|
|
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.script
|
|
from tvm import relax
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def test_vm_builtin_lower_mem_alloc_storage():
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor:
|
|
R.func_attr({"relax.force_pure": True})
|
|
m, n = T.int64(), T.int64()
|
|
|
|
storage = R.memory.alloc_storage(R.shape([m * n * 4]), 0, "global", "uint8")
|
|
alloc = R.memory.alloc_tensor(storage, 0, R.shape([m, n]), "float32")
|
|
_ = R.call_packed(
|
|
"test.op.identity", x, alloc, ty_args=(R.Tensor(ndim=2, dtype="float32"))
|
|
)
|
|
gv0 = alloc
|
|
return gv0
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor:
|
|
# we expected RemovePurityChecking to have been called first
|
|
R.func_attr({"relax.force_pure": True})
|
|
m, n = T.int64(), T.int64()
|
|
|
|
storage = R.vm.alloc_storage(R.shape([m * n * 4]), R.prim_value(0), "uint8", "global")
|
|
alloc = R.vm.alloc_tensor(storage, R.prim_value(0), R.shape([m, n]), "float32")
|
|
|
|
_ = R.call_packed(
|
|
"test.op.identity", x, alloc, ty_args=(R.Tensor(ndim=2, dtype="float32"))
|
|
)
|
|
gv0 = alloc
|
|
return gv0
|
|
|
|
After = relax.transform.LowerRuntimeBuiltin()(Before)
|
|
tvm.ir.assert_structural_equal(Expected, After)
|
|
|
|
|
|
def test_vm_builtin_alloc_tensor_raises_error():
|
|
"""R.builtin.alloc_tensor should be handled earlier"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor:
|
|
R.func_attr({"relax.force_pure": True})
|
|
m, n = T.int64(), T.int64()
|
|
|
|
alloc = R.builtin.alloc_tensor(R.shape([m, n]), runtime_device_index=0, dtype="float32")
|
|
_ = R.call_packed(
|
|
"test.op.identity", x, alloc, ty_args=(R.Tensor(ndim=2, dtype="float32"))
|
|
)
|
|
gv0 = alloc
|
|
return gv0
|
|
|
|
with pytest.raises(RuntimeError):
|
|
relax.transform.LowerRuntimeBuiltin()(Before)
|
|
|
|
|
|
def test_vm_reshape_may_be_var():
|
|
"""R.reshape does not require an in-line R.shape"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(A: R.Tensor([16], "float32"), shape: R.Shape):
|
|
R.func_attr({"relax.force_pure": True})
|
|
reshape = R.reshape(A, shape)
|
|
return reshape
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(A: R.Tensor([16], "float32"), shape: R.Shape):
|
|
R.func_attr({"relax.force_pure": True})
|
|
reshape = R.call_packed(
|
|
"vm.builtin.reshape",
|
|
A,
|
|
shape,
|
|
ty_args=R.Tensor(shape, dtype="float32"),
|
|
)
|
|
return reshape
|
|
|
|
After = relax.transform.VMBuiltinLower()(Before)
|
|
|
|
tvm.ir.assert_structural_equal(Expected, After)
|
|
|
|
|
|
def test_vm_reshape_using_tensor_to_shape():
|
|
"""Shape argument of R.reshape may come from tensor_to_shape"""
|
|
|
|
@I.ir_module
|
|
class Before:
|
|
@R.function
|
|
def main(A: R.Tensor([16], "float32"), shape_tensor: R.Tensor([2], "int64")):
|
|
R.func_attr({"relax.force_pure": True})
|
|
shape = R.tensor_to_shape(shape_tensor)
|
|
reshape = R.reshape(A, shape)
|
|
return reshape
|
|
|
|
@I.ir_module
|
|
class Expected:
|
|
@R.function
|
|
def main(A: R.Tensor([16], "float32"), shape_tensor: R.Tensor([2], "int64")):
|
|
R.func_attr({"relax.force_pure": True})
|
|
|
|
shape = R.call_packed(
|
|
"vm.builtin.tensor_to_shape",
|
|
shape_tensor,
|
|
ty_args=R.Shape(ndim=2),
|
|
)
|
|
reshape = R.call_packed(
|
|
"vm.builtin.reshape",
|
|
A,
|
|
shape,
|
|
ty_args=R.Tensor(shape, dtype="float32"),
|
|
)
|
|
return reshape
|
|
|
|
After = relax.transform.VMBuiltinLower()(Before)
|
|
|
|
tvm.ir.assert_structural_equal(Expected, After)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|