python/mxnet/tvmop.py

# 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
"""Init tvm ops."""
from .runtime import Features

if Features().is_enabled("TVM_OP"):
    import json
    import logging

    from ._ctypes.space import _set_tvm_op_config
    from .base import check_call, _LIB, c_str
    from .space import ConfigSpaces
    from .libinfo import find_lib_path, find_conf_path

    _LIB_TVM_OP = find_lib_path("libtvmop")
    check_call(_LIB.MXLoadTVMOp(c_str(_LIB_TVM_OP[0])))

    # op sch config
    try:
        _CONF_TVM_OP = find_conf_path("tvmop")
    except RuntimeError as e:
        logging.warning("TVM config file missing, falling back to default schedule", exc_info=True)
    else:
        logging.info("TVM op config has been loaded")
        with open(_CONF_TVM_OP[0], "r") as f:
            ret = ConfigSpaces.from_json_dict(json.load(f))
        _set_tvm_op_config(ret)