From 44838c495ccaffadac673f32b3d7efca542f8e63 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Fri, 17 Feb 2023 21:40:39 -0500 Subject: [PATCH] [Unity][TVMScript] Move tir/relax import in script out of __init__.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Prior to this PR, `python/tvm/script/__init__.py` imports both tir and relax submodules. This leads to the phenomenum that when people does ```python from tvm.script import tir as T ``` , the relax submodule will be implicitly visited by `__init__.py` as well. Since TIR does not rely on Relax, it is good not to import both of them at the same time. (This can prevent cyclic imports sometimes.) This PR does this decoupling by introducing two files * `python/tvm/script/relax.py` * `python/tvm/script/tir.py` and removing the imports from `python/tvm/script/__init__.py` and `python/tvm/script/parser/__init__.py`. With this change, we force people to manually do `from tvm.script import tir` and `from tvm.script import relax` to use TVMScript parser, which is right our conventional way. --- To elabore a bit more on the cyclic import: With #14029, submodule `relax.transform` will import topi to register the legalization functions for each operators into the op registry. Module topi imports the TVMScript parser of TIR, in `python/tvm/topi/arm_cpu.py`. As mentioned above, the previous behavior is when importing TIR parser, the submodule `tvm.script.parser.relax` is also visited and initialized, and thereby brings the risk of cyclic import. Since the TOPI side only needs TIR parser, we don’t need to import the Relax parser. By decoupling the import of TIR and Relax in module `tvm.script` can prevent the cyclic dependency. --- python/tvm/script/__init__.py | 2 -- python/tvm/script/parser/__init__.py | 6 ++---- python/tvm/script/parser/tir/__init__.py | 2 +- python/tvm/script/relax.py | 18 ++++++++++++++++++ python/tvm/script/tir.py | 18 ++++++++++++++++++ 5 files changed, 39 insertions(+), 7 deletions(-) create mode 100644 python/tvm/script/relax.py create mode 100644 python/tvm/script/tir.py diff --git a/python/tvm/script/__init__.py b/python/tvm/script/__init__.py index 6d92c68367b3..f5ee692cbb8f 100644 --- a/python/tvm/script/__init__.py +++ b/python/tvm/script/__init__.py @@ -17,5 +17,3 @@ """TVM Script APIs of TVM Python Package""" from .parser import ir, ir_module from .parser import parse as from_source -from .parser import tir -from .parser import relax diff --git a/python/tvm/script/parser/__init__.py b/python/tvm/script/parser/__init__.py index 678297799e6d..ba7f085c08a4 100644 --- a/python/tvm/script/parser/__init__.py +++ b/python/tvm/script/parser/__init__.py @@ -13,10 +13,8 @@ # "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 Licens. +# under the License. """The parser""" -from . import _core, ir, tir, relax +from . import _core, ir from ._core import parse from .ir import ir_module -from .tir import prim_func -from .relax import function diff --git a/python/tvm/script/parser/tir/__init__.py b/python/tvm/script/parser/tir/__init__.py index ad16821a89a3..e44b6b521b27 100644 --- a/python/tvm/script/parser/tir/__init__.py +++ b/python/tvm/script/parser/tir/__init__.py @@ -32,4 +32,4 @@ else: from .entry import prim_func -__all__ = _tir.__all__ + ["Buffer", "Ptr", "prim_func"] +__all__ = _tir.__all__ + ["Buffer", "Ptr", "bool", "prim_func"] diff --git a/python/tvm/script/relax.py b/python/tvm/script/relax.py new file mode 100644 index 000000000000..2301463059e3 --- /dev/null +++ b/python/tvm/script/relax.py @@ -0,0 +1,18 @@ +# 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. +"""TVM Script APIs of TVM Python Package for Relax""" +from .parser.relax import * # pylint: disable=redefined-builtin,unused-wildcard-import,wildcard-import diff --git a/python/tvm/script/tir.py b/python/tvm/script/tir.py new file mode 100644 index 000000000000..49f3ecd42c50 --- /dev/null +++ b/python/tvm/script/tir.py @@ -0,0 +1,18 @@ +# 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. +"""TVM Script APIs of TVM Python Package for TIR""" +from .parser.tir import * # pylint: disable=redefined-builtin,unused-wildcard-import,wildcard-import