chore: import upstream snapshot with attribution

python/tvm/__init__.py

@@ -0,0 +1,126 @@
+# isort: skip_file
+# 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.
+# pylint: disable=redefined-builtin, wildcard-import
+"""TVM: Open Deep Learning Compiler Stack."""
+
+import multiprocessing
+import sys
+import os
+
+# ffi module must load first
+from tvm_ffi import register_object, register_global_func, get_global_func
+
+# top-level alias
+from .libinfo import __version__
+from .base import _RUNTIME_ONLY
+
+# tvm.runtime
+from .runtime import Object
+from .runtime._tensor import device, cpu, cuda, opencl, vulkan, metal
+from .runtime._tensor import vpi, rocm, ext_dev, hexagon
+from .runtime import DataType, DataTypeCode
+
+# tvm.error
+from . import error
+
+# tvm.ir
+from .ir import IRModule
+from .ir import transform
+from .ir import instrument
+from . import ir
+
+# tvm.script — must be imported before any dialect package so that
+# tvm.script.register_dialect is reachable when dialect __init__.py files run.
+from . import script
+
+# tvm.tirx — registers itself via tvm.script.register_dialect in its __init__
+from . import tirx
+
+# tvm.backend — owns backend Python load hooks
+from . import backend
+
+# tvm.target
+from . import target
+
+# tvm.te
+from . import te
+
+# tvm.driver
+from .driver import build, compile
+
+# others
+from . import arith
+
+# support infra
+from . import support
+
+# Side-effect imports: register CUDA/ROCm FFI callbacks at TVM startup
+from .support import rocm as _rocm, nvcc as _nvcc
+
+# Relax contain modules that are only available in compiler package
+# Do not import them if TVM is built with runtime only
+if not _RUNTIME_ONLY:
+    # tvm.relax — registers itself via tvm.script.register_dialect in its __init__
+    from . import relax
+
+# NOTE: This file should be python2 compatible so we can
+# raise proper error message when user run the package using
+# an older version of the python
+
+
+def _should_print_backtrace():
+    in_pytest = "PYTEST_CURRENT_TEST" in os.environ
+    tvm_backtrace = os.environ.get("TVM_BACKTRACE", "0")
+
+    try:
+        tvm_backtrace = bool(int(tvm_backtrace))
+    except ValueError:
+        raise ValueError(f"invalid value for TVM_BACKTRACE {tvm_backtrace}, please set to 0 or 1.")
+
+    return in_pytest or tvm_backtrace
+
+
+def tvm_wrap_excepthook(exception_hook):
+    """Wrap given excepthook with TVM additional work."""
+
+    def wrapper(exctype, value, trbk):
+        """Clean subprocesses when TVM is interrupted."""
+        if exctype is error.DiagnosticError and not _should_print_backtrace():
+            # TODO(@jroesch): consider moving to C++?
+            print("note: run with `TVM_BACKTRACE=1` environment variable to display a backtrace.")
+        else:
+            exception_hook(exctype, value, trbk)
+
+        if hasattr(multiprocessing, "active_children"):
+            # pylint: disable=not-callable
+            for p in multiprocessing.active_children():
+                p.terminate()
+
+    return wrapper
+
+
+sys.excepthook = tvm_wrap_excepthook(sys.excepthook)
+
+# Autoload loads built-in and out-of-tree backends. Out-of-tree extensions opt
+# into being loaded automatically at ``import tvm`` time by declaring an entry
+# point in the ``tvm.backends`` group:
+# [project.entry-points."tvm.backends"] tvm_foo = "tvm_foo:_autoload".
+# Autoload can be disabled via ``TVM_DEVICE_BACKEND_AUTOLOAD=0``.
+from .backend._autoload_backends import _autoload_backends
+
+_autoload_backends()