blob: 6f9fb7c5fd61a51037dbcee3bf7b5bac33103194 [file] [log] [blame]
# 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
"""TVM Operator compile entry point"""
import tvm
from tvm import autotvm
import os
import argparse
import re
import json
import logging
import sys
import subprocess
from tvmop.opdef import __OP_DEF__
from tvmop.space import ConfigSpaces, ConfigSpace
from tvm.autotvm.measure.measure_methods import set_cuda_target_arch
logging.basicConfig(level=logging.INFO)
def create_shared(output,
objects,
options=None,
cc="g++"):
"""Create shared library.
Parameters
----------
output : str
The target shared library.
objects : List[str]
List of object files.
options : List[str]
The list of additional options string.
cc : Optional[str]
The compiler command.
"""
if sys.platform == "darwin" or sys.platform.startswith("linux"):
_linux_compile(output, objects, options, cc)
# TODO(yzhliu): elif sys.platform == "win32":
else:
raise ValueError("Unsupported platform")
def _linux_compile(output, objects, options, compile_cmd="g++"):
cmd = [compile_cmd]
if output.endswith(".so") or output.endswith(".dylib"):
cmd += ["-shared", "-fPIC"]
if sys.platform == "darwin":
cmd += ["-undefined", "dynamic_lookup"]
elif output.endswith(".obj"):
cmd += ["-c"]
cmd += ["-o", output]
if isinstance(objects, str):
cmd += [objects]
else:
cmd += objects
if options:
cmd += options
proc = subprocess.Popen(
cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
(out, _) = proc.communicate()
if proc.returncode != 0:
msg = "Compilation error:\n"
msg += str(out)
raise RuntimeError(msg)
def get_target(device):
if device == "cpu":
return "llvm"
elif device == "cuda" or device == "gpu":
return "cuda"
assert False, "Unknown device " + device
def get_cuda_arch(arch):
if arch is None:
return None
if not isinstance(arch, str):
raise TypeError('Expecting parameter arch as a str, while got a {}'.format(str(type(arch))))
if len(arch) == 0:
return None
# an example of arch string,
# -gencode arch=compute_30,code=sm_30 -gencode arch=compute_35,code=sm_35
# -gencode;arch=compute_75,code=[sm_75,compute_75] --fatbin-options -compress-all
archs = []
flags = arch.replace("-gencode;", "-gencode ").split()
for flag in flags:
if flag.startswith('-gencode') or flag.startswith('arch='):
archs.append(flag)
return archs
if __name__ == "__main__":
import sys
sys.path.append(os.path.dirname(sys.path[0]))
parser = argparse.ArgumentParser(description="Generate tvm operators")
parser.add_argument("-o", action="store", required=True, dest="target_path",
help="Target path which stores compiled library")
parser.add_argument("-L", action="store", default=None, dest="ld_path",
help="library link path")
parser.add_argument('--cuda-arch', type=str, default=None, dest='cuda_arch',
help='The cuda arch for compiling kernels for')
parser.add_argument("--config", action="store", required=True, dest="config_path",
help="Path which stores the config file")
arguments = parser.parse_args()
mod_llvm = tvm.IRModule({})
mod_cuda = tvm.IRModule({})
has_cuda = False
# TODO: attach instruction features to the library, e.g., avx-512, etc.
for operator_def in __OP_DEF__:
for sch, args, name in operator_def.invoke_all():
name = operator_def.get_op_name(name, args)
if tvm.runtime.module.enabled(get_target(operator_def.target)):
func_lower = tvm.lower(sch, args,
name=name,
binds=operator_def.get_binds(args))
if operator_def.target == "cpu":
mod = mod_llvm.update(func_lower)
else:
has_cuda = True
mod_cuda.update(func_lower)
lowered_funcs = {get_target("cpu"): mod_llvm}
if has_cuda > 0:
lowered_funcs[get_target("cuda")] = mod_cuda
cuda_arch = get_cuda_arch(arguments.cuda_arch)
if cuda_arch is None:
logging.info('No cuda arch specified. TVM will try to detect it from the build platform.')
else:
logging.info('Cuda arch {} set for compiling TVM operator kernels.'.format(cuda_arch))
set_cuda_target_arch(cuda_arch)
func_binary = tvm.build(lowered_funcs, name="tvmop")
# we create libtvmop.o first, which gives us chance to link tvm_runtime together with the libtvmop
# to allow mxnet find external helper functions in libtvm_runtime
func_binary.save(arguments.target_path + "/libtvmop.o")
if len(func_binary.imported_modules):
func_binary.imported_modules[0].save(arguments.target_path + "/libtvmop.cubin")
ld_path = arguments.target_path if arguments.ld_path is None else arguments.ld_path
create_shared(arguments.target_path + "/libtvmop.so",
arguments.target_path + "/libtvmop.o",
options=["-L", ld_path, "-ltvm_runtime"])
config_spaces = ConfigSpaces()
for operator_def in __OP_DEF__:
for config_space, name in operator_def.get_config_spaces():
config_spaces[name] = ConfigSpace.from_tvm(config_space)
with open(arguments.config_path, "w") as f:
json.dump(config_spaces.to_json_dict(), f)