1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements.  See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership.  The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License.  You may obtain a copy of the License at
8#
9#   http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied.  See the License for the
15# specific language governing permissions and limitations
16# under the License.
17# pylint: disable=redefined-builtin, wildcard-import
18"""TVM: Low level DSL/IR stack for tensor computation."""
19from __future__ import absolute_import as _abs
20
21import multiprocessing
22import sys
23import traceback
24
25from . import _pyversion
26
27from . import tensor
28from . import arith
29from . import expr
30from . import stmt
31from . import make
32from . import ir_pass
33from . import codegen
34from . import container
35from . import schedule
36from . import module
37from . import node
38from . import attrs
39from . import ir_builder
40from . import target
41from . import generic
42from . import hybrid
43from . import testing
44from . import error
45from . import datatype
46
47from . import ndarray as nd
48from .ndarray import context, cpu, gpu, opencl, cl, vulkan, metal, mtl
49from .ndarray import vpi, rocm, opengl, ext_dev, micro_dev
50
51from ._ffi.runtime_ctypes import TypeCode, TVMType
52from ._ffi.ndarray import TVMContext
53from ._ffi.function import Function
54from ._ffi.base import TVMError, __version__
55from .api import *
56from .intrin import *
57from .tensor_intrin import decl_tensor_intrin
58from .node import register_node
59from .ndarray import register_extension
60from .schedule import create_schedule
61from .build_module import build, lower, build_config
62from .tag import tag_scope
63
64# Contrib initializers
65from .contrib import rocm as _rocm, nvcc as _nvcc, sdaccel as _sdaccel
66
67# Clean subprocesses when TVM is interrupted
68def tvm_excepthook(exctype, value, trbk):
69    print('\n'.join(traceback.format_exception(exctype, value, trbk)))
70    if hasattr(multiprocessing, 'active_children'):
71        # pylint: disable=not-callable
72        for p in multiprocessing.active_children():
73            p.terminate()
74
75sys.excepthook = tvm_excepthook
76