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