wehub-resource-sync
62 lines
2.3 KiB
Python
62 lines
2.3 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.
|
|
# coding: utf-8
|
|
# pylint: disable=invalid-name, import-outside-toplevel
|
|
"""Base library for TVM."""
|
|
|
|
import os
|
|
|
|
from tvm_ffi.libinfo import load_lib_ctypes
|
|
|
|
from . import libinfo
|
|
|
|
# ----------------------------
|
|
# library loading
|
|
# ----------------------------
|
|
|
|
# Whether only the runtime library is loaded (runtime-only wheel, or
|
|
# ``TVM_USE_RUNTIME_LIB=1``). Set during library loading below.
|
|
_RUNTIME_ONLY = os.environ.get("TVM_USE_RUNTIME_LIB") == "1"
|
|
|
|
# Handles of the core libraries actually loaded, keyed by basename
|
|
# (e.g. ``{"tvm_runtime": <CDLL>, "tvm_compiler": <CDLL>}``). Downstream /
|
|
# autoloaded extensions can inspect this to skip duplicate libraries
|
|
# (``"tvm_runtime" in _LOADED_LIBS``) and obtain the loaded handle.
|
|
_LOADED_LIBS = {}
|
|
|
|
|
|
# runtime is loaded RTLD_GLOBAL to expose its symbols to subsequent loads;
|
|
# compiler is loaded RTLD_LOCAL.
|
|
_LOADED_LIBS["tvm_runtime"] = load_lib_ctypes(
|
|
"tvm", "tvm_runtime", "RTLD_GLOBAL", extra_lib_paths=libinfo.package_lib_paths()
|
|
)
|
|
|
|
if not _RUNTIME_ONLY:
|
|
try:
|
|
_LOADED_LIBS["tvm_compiler"] = load_lib_ctypes(
|
|
"tvm", "tvm_compiler", "RTLD_LOCAL", extra_lib_paths=libinfo.package_lib_paths()
|
|
)
|
|
except (RuntimeError, OSError):
|
|
# Compiler lib not present, or present but unloadable (missing LLVM
|
|
# deps / linker issues) — fall back to runtime-only mode.
|
|
_RUNTIME_ONLY = True
|
|
|
|
if _RUNTIME_ONLY:
|
|
from tvm_ffi import registry as _tvm_ffi_registry
|
|
|
|
_tvm_ffi_registry._SKIP_UNKNOWN_OBJECTS = True
|