diff --git a/3rdparty/dlpack b/3rdparty/dlpack index 10892ac964f1..b90e93907206 160000 --- a/3rdparty/dlpack +++ b/3rdparty/dlpack @@ -1 +1 @@ -Subproject commit 10892ac964f1af7c81aae145cd3fab78bbccd297 +Subproject commit b90e939072066c160b18ea1e7156537b8d3710f6 diff --git a/3rdparty/dmlc-core b/3rdparty/dmlc-core index 3943914eed66..f1ff6cc117f4 160000 --- a/3rdparty/dmlc-core +++ b/3rdparty/dmlc-core @@ -1 +1 @@ -Subproject commit 3943914eed66470bd010df581e29e4dca4f7df6f +Subproject commit f1ff6cc117f4e95169a9f62be549c8fe3e15c20f diff --git a/3rdparty/tvm b/3rdparty/tvm index 21935dcbf56a..afd4b3e44509 160000 --- a/3rdparty/tvm +++ b/3rdparty/tvm @@ -1 +1 @@ -Subproject commit 21935dcbf56ad3bd66ebff9891a6bc3865b8106d +Subproject commit afd4b3e4450984358e9d79a7e8e578483cb7b017 diff --git a/CMakeLists.txt b/CMakeLists.txt index 0148ac302d54..42463cd227c6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -43,6 +43,7 @@ mxnet_option(USE_MXNET_LIB_NAMING "Use MXNet library naming conventions." ON) mxnet_option(USE_GPROF "Compile with gprof (profiling) flag" OFF) mxnet_option(USE_CXX14_IF_AVAILABLE "Build with C++14 if the compiler supports it" OFF) mxnet_option(USE_VTUNE "Enable use of Intel Amplifier XE (VTune)" OFF) # one could set VTUNE_ROOT for search path +mxnet_option(USE_TVM_OP "Enable use of TVM operator build system." OFF) mxnet_option(ENABLE_CUDA_RTC "Build with CUDA runtime compilation support" ON) mxnet_option(BUILD_CPP_EXAMPLES "Build cpp examples" ON) mxnet_option(INSTALL_EXAMPLES "Install the example source files." OFF) @@ -734,6 +735,28 @@ if(USE_DIST_KVSTORE) list(APPEND mxnet_LINKER_LIBS ${pslite_LINKER_LIBS}) endif() +if(USE_TVM_OP) + add_definitions(-DMXNET_USE_TVM_OP=1) + list(APPEND mxnet_LINKER_LIBS ${CMAKE_CURRENT_BINARY_DIR}/3rdparty/tvm/libtvm_runtime.so) + include(cmake/BuildTVM.cmake) + add_subdirectory("3rdparty/tvm") + + if(NOT Python3_EXECUTABLE) + find_package(PythonInterp 3 REQUIRED) + set(Python3_EXECUTABLE ${PYTHON_EXECUTABLE} CACHE FILEPATH "Path to the python3 executable") + if(NOT Python3_EXECUTABLE) + message(FATAL_ERROR "No python3 interpreter found to build TVM operators") + endif() + endif() + + add_custom_command(TARGET mxnet POST_BUILD + COMMAND ${CMAKE_COMMAND} -E env + PYTHONPATH="${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm/python:${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm/topi/python:${CMAKE_CURRENT_SOURCE_DIR}/contrib" + LD_LIBRARY_PATH="${CMAKE_CURRENT_BINARY_DIR}/3rdparty/tvm/build" + ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/contrib/tvmop/compile.py -o${CMAKE_CURRENT_BINARY_DIR}/libtvmop.so + ) +endif() + target_link_libraries(mxnet PUBLIC ${mxnet_LINKER_LIBS}) if(USE_PLUGINS_WARPCTC) diff --git a/Makefile b/Makefile index a7cf8a56a433..b11e2969591f 100644 --- a/Makefile +++ b/Makefile @@ -52,6 +52,14 @@ ifndef AMALGAMATION_PATH AMALGAMATION_PATH = $(ROOTDIR)/amalgamation endif +ifndef TVM_PATH + TVM_PATH = $(TPARTYDIR)/tvm +endif + +ifndef LLVM_PATH + LLVM_PATH = $(TVM_PATH)/build/llvm +endif + ifneq ($(USE_OPENMP), 1) export NO_OPENMP = 1 endif @@ -589,6 +597,35 @@ $(DMLC_CORE)/libdmlc.a: DMLCCORE DMLCCORE: + cd $(DMLC_CORE); $(MAKE) libdmlc.a USE_SSE=$(USE_SSE) config=$(ROOTDIR)/$(config); cd $(ROOTDIR) +ifeq ($(USE_TVM_OP), 1) +LIB_DEP += lib/libtvm_runtime.so lib/libtvmop.so +CFLAGS += -I$(TVM_PATH)/include -DMXNET_USE_TVM_OP=1 +LDFLAGS += -L$(TVM_PATH)/build -ltvm_runtime + +TVM_USE_CUDA := OFF +ifeq ($(USE_CUDA), 1) + TVM_USE_CUDA := ON + ifneq ($(USE_CUDA_PATH), NONE) + TVM_USE_CUDA := $(USE_CUDA_PATH) + endif +endif +lib/libtvm_runtime.so: + echo "Compile TVM" + [ -e $(LLVM_PATH)/bin/llvm-config ] || sh $(ROOTDIR)/contrib/tvmop/prepare_tvm.sh; \ + cd $(TVM_PATH)/build; \ + cmake -DUSE_LLVM="$(LLVM_PATH)/bin/llvm-config" \ + -DUSE_SORT=OFF -DUSE_CUDA=$(TVM_USE_CUDA) -DUSE_CUDNN=OFF ..; \ + $(MAKE) VERBOSE=1; \ + cp $(TVM_PATH)/build/libtvm_runtime.so $(ROOTDIR)/lib/libtvm_runtime.so; \ + cd $(ROOTDIR) + +lib/libtvmop.so: lib/libtvm_runtime.so $(wildcard contrib/tvmop/*/*.py contrib/tvmop/*.py) + echo "Compile TVM operators" + PYTHONPATH=$(TVM_PATH)/python:$(TVM_PATH)/topi/python:$(ROOTDIR)/contrib:$PYTHONPATH \ + LD_LIBRARY_PATH=lib \ + python3 $(ROOTDIR)/contrib/tvmop/compile.py -o $(ROOTDIR)/lib/libtvmop.so +endif + NNVM_INC = $(wildcard $(NNVM_PATH)/include/*/*.h) NNVM_SRC = $(wildcard $(NNVM_PATH)/src/*/*/*.cc $(NNVM_PATH)/src/*/*.cc $(NNVM_PATH)/src/*.cc) $(NNVM_PATH)/lib/libnnvm.a: $(NNVM_INC) $(NNVM_SRC) @@ -726,6 +763,7 @@ clean: rclean cyclean $(EXTRA_PACKAGES_CLEAN) cd $(DMLC_CORE); $(MAKE) clean; cd - cd $(PS_PATH); $(MAKE) clean; cd - cd $(NNVM_PATH); $(MAKE) clean; cd - + cd $(TVM_PATH); $(MAKE) clean; cd - cd $(AMALGAMATION_PATH); $(MAKE) clean; cd - $(RM) -r $(patsubst %, %/*.d, $(EXTRA_OPERATORS)) $(patsubst %, %/*/*.d, $(EXTRA_OPERATORS)) $(RM) -r $(patsubst %, %/*.o, $(EXTRA_OPERATORS)) $(patsubst %, %/*/*.o, $(EXTRA_OPERATORS)) diff --git a/cmake/BuildTVM.cmake b/cmake/BuildTVM.cmake new file mode 100644 index 000000000000..ad8517cafe26 --- /dev/null +++ b/cmake/BuildTVM.cmake @@ -0,0 +1,135 @@ +# 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. + +message(STATUS "Prepare external packages for TVM...") +execute_process(COMMAND "${CMAKE_CURRENT_SOURCE_DIR}/contrib/tvmop/prepare_tvm.sh") + +# Whether enable ROCM runtime +# +# Possible values: +# - ON: enable ROCM with cmake's auto search +# - OFF: disable ROCM +# - /path/to/rocm: use specific path to rocm +set(USE_ROCM OFF) + +# Whether enable SDAccel runtime +set(USE_SDACCEL OFF) + +# Whether enable Intel FPGA SDK for OpenCL (AOCL) runtime +set(USE_AOCL OFF) + +# Whether enable OpenCL runtime +set(USE_OPENCL OFF) + +# Whether enable Metal runtime +set(USE_METAL OFF) + +# Whether enable Vulkan runtime +# +# Possible values: +# - ON: enable Vulkan with cmake's auto search +# - OFF: disable vulkan +# - /path/to/vulkan-sdk: use specific path to vulkan-sdk +set(USE_VULKAN OFF) + +# Whether enable OpenGL runtime +set(USE_OPENGL OFF) + +# Whether to enable SGX runtime +# +# Possible values for USE_SGX: +# - /path/to/sgxsdk: path to Intel SGX SDK +# - OFF: disable SGX +# +# SGX_MODE := HW|SIM +set(USE_SGX OFF) +set(SGX_MODE "SIM") +set(RUST_SGX_SDK "/path/to/rust-sgx-sdk") + +# Whether enable RPC runtime +set(USE_RPC ON) + +# Whether embed stackvm into the runtime +set(USE_STACKVM_RUNTIME OFF) + +# Whether enable tiny embedded graph runtime. +set(USE_GRAPH_RUNTIME ON) + +# Whether enable additional graph debug functions +set(USE_GRAPH_RUNTIME_DEBUG OFF) + +# Whether build with LLVM support +# Requires LLVM version >= 4.0 +# +# Possible values: +# - ON: enable llvm with cmake's find search +# - OFF: disable llvm +# - /path/to/llvm-config: enable specific LLVM when multiple llvm-dev is available. +set(USE_LLVM "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm/build/llvm/bin/llvm-config") + +#--------------------------------------------- +# Contrib libraries +#--------------------------------------------- +# Whether use BLAS, choices: openblas, mkl, atlas, apple +set(USE_BLAS none) + +# /path/to/mkl: mkl root path when use mkl blas library +# set(USE_MKL_PATH /opt/intel/mkl) for UNIX +# set(USE_MKL_PATH ../IntelSWTools/compilers_and_libraries_2018/windows/mkl) for WIN32 +set(USE_MKL_PATH none) + +# Whether use contrib.random in runtime +set(USE_RANDOM OFF) + +# Whether use NNPack +set(USE_NNPACK OFF) + +# Whether use CuDNN +if(USE_CUDNN AND USE_CUDA) + detect_cuDNN() + if(HAVE_CUDNN) + set(USE_CUDNN ON) + else() + set(USE_CUDNN OFF) + endif() +else() + set(USE_CUDNN OFF) +endif() + +# Whether use cuBLAS +set(USE_CUBLAS OFF) + +# Whether use MIOpen +set(USE_MIOPEN OFF) + +# Whether use MPS +set(USE_MPS OFF) + +# Whether use rocBlas +set(USE_ROCBLAS OFF) + +# Whether use contrib sort +set(USE_SORT OFF) + +# Build ANTLR parser for Relay text format +set(USE_ANTLR OFF) + +# Build TSIM for VTA +set(USE_VTA_TSIM OFF) + +# Whether use Relay debug mode +set(USE_RELAY_DEBUG OFF) diff --git a/contrib/tvmop/__init__.py b/contrib/tvmop/__init__.py new file mode 100644 index 000000000000..31189d499b5a --- /dev/null +++ b/contrib/tvmop/__init__.py @@ -0,0 +1,22 @@ +# 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 +from .opdef import defop +from .utils import AllTypes, RealTypes + +from . import basic diff --git a/contrib/tvmop/basic/__init__.py b/contrib/tvmop/basic/__init__.py new file mode 100644 index 000000000000..fc0fa720009f --- /dev/null +++ b/contrib/tvmop/basic/__init__.py @@ -0,0 +1,19 @@ +# 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 +from . import ufunc diff --git a/contrib/tvmop/basic/ufunc.py b/contrib/tvmop/basic/ufunc.py new file mode 100644 index 000000000000..0419e5fd2ca9 --- /dev/null +++ b/contrib/tvmop/basic/ufunc.py @@ -0,0 +1,50 @@ +# 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 +import tvm +from .. import defop, AllTypes + +def compute_add(dtype, ndim): + A = tvm.placeholder([tvm.var() for _ in range(ndim)], name='A', dtype=dtype) + B = tvm.placeholder([tvm.var() for _ in range(ndim)], name='B', dtype=dtype) + C = tvm.compute([tvm.var() for _ in range(ndim)], + lambda *index: A[index] + B[index], name='C') + s = tvm.create_schedule(C.op) + return s, A, B, C + +@defop(name="vadd", target="cpu", auto_broadcast=True, + dtype=AllTypes, ndim=list(range(1, 6))) +def vadd(dtype, ndim): + s, A, B, C = compute_add(dtype, ndim) + axes = [axis for axis in C.op.axis] + fused = s[C].fuse(*axes) + s[C].parallel(fused) + + return s, [A, B, C] + +@defop(name="cuda_vadd", target="cuda", auto_broadcast=True, + dtype=["float32", "float64"], ndim=list(range(1, 6))) +def vadd_gpu(dtype, ndim): + s, A, B, C = compute_add(dtype, ndim) + s = tvm.create_schedule(C.op) + axes = [axis for axis in C.op.axis] + fused = s[C].fuse(*axes) + bx, tx = s[C].split(fused, factor=64) + s[C].bind(bx, tvm.thread_axis("blockIdx.x")) + s[C].bind(tx, tvm.thread_axis("threadIdx.x")) + return s, [A, B, C] diff --git a/contrib/tvmop/compile.py b/contrib/tvmop/compile.py new file mode 100644 index 000000000000..94274fe4142a --- /dev/null +++ b/contrib/tvmop/compile.py @@ -0,0 +1,59 @@ +# 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 + +import os +import argparse +from tvmop.opdef import __OP_DEF__ + +def get_target(device): + if device == "cpu": + return "llvm" + elif device == "cuda" or device == "gpu": + return "cuda" + assert False, "Unknown device " + device + + +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") + arguments = parser.parse_args() + + func_list_llvm = [] + func_list_cuda = [] + + # TODO: attach instruction features to the library, e.g., avx-512, etc. + for operator_def in __OP_DEF__: + for sch, args in operator_def.invoke_all(): + if tvm.module.enabled(get_target(operator_def.target)): + func_list = func_list_llvm if operator_def.target == "cpu" else func_list_cuda + func_lower = tvm.lower(sch, args, + name=operator_def.get_op_name(args), + binds=operator_def.get_binds(args)) + func_list.append(func_lower) + + lowered_funcs = {get_target("cpu") : func_list_llvm} + if len(func_list_cuda) > 0: + lowered_funcs[get_target("cuda")] = func_list_cuda + func_binary = tvm.build(lowered_funcs, name="tvmop") + func_binary.export_library(arguments.target_path) diff --git a/contrib/tvmop/opdef.py b/contrib/tvmop/opdef.py new file mode 100644 index 000000000000..c65824588047 --- /dev/null +++ b/contrib/tvmop/opdef.py @@ -0,0 +1,111 @@ +# 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 +import tvm +from itertools import product + +__OP_DEF__ = [] + +class OpDef: + """Specify the properties of an operator and + construct the value combination of the arguments + e.g., ldtype=["float32", "int32"], rdtype=["float16", "int16"], + then the argument combination is + [ + {"ldtype": "float32", "rdtype": "float16"}, + {"ldtype": "float32", "rdtype": "int16"}, + {"ldtype": "int32", "rdtype": "float16"}, + {"ldtype": "int32", "rdtype": "int16"}, + ] + + Parameters + ---------- + func : function + The function to define the operator (in tvm compute and schedule). + It will get the argument combination extracted by this class. + name : str + function name. + target : str + {"cpu", "gpu", "cuda"} + auto_broadcast : bool + auto_broadcast=True allows one to implement broadcast computation + without considering whether dimension size equals to one. + TVM maps buffer[i][j][k] -> buffer[i][0][k] if dimension i's shape equals 1. + """ + def __init__(self, func, name, target, auto_broadcast, **kwargs): + # construct the value combination of the arguments + # e.g., ldtype=["float32", "int32"], rdtype=["float16", "int16"] + # arg_combination = [ + # {"ldtype": "float32", "rdtype": "float16"}, + # {"ldtype": "float32", "rdtype": "int16"}, + # {"ldtype": "int32", "rdtype": "float16"}, + # {"ldtype": "int32", "rdtype": "int16"}, + # ] + args = [k for k in kwargs] + values = [kwargs[k] if isinstance(kwargs[k], (list, tuple)) else [kwargs[k]] + for k in args] + cart_product = product(*values) + self.arg_combination = [{k: v for k, v in zip(args, comb_values)} + for comb_values in cart_product] + self.func = func + self.name = name + self.target = target + self.auto_broadcast = auto_broadcast + + def __call__(self, *args, **kwargs): + return self.func(*args, **kwargs) + + def invoke_all(self): + for each_kwargs in self.arg_combination: + yield self.func(**each_kwargs) + + def get_op_name(self, args): + return self.name + ''.join(["%s_%d" % (arg.dtype, len(arg.shape)) for arg in args]) + + def get_binds(self, args): + if self.auto_broadcast: + return {arg: tvm.decl_buffer(arg.shape, arg.dtype, buffer_type="auto_broadcast") + for arg in args} + return None + + +def defop(name, target=None, auto_broadcast=False, **kwargs): + """Decorator to define a tvm operator. + Parameters + ---------- + name : str + function name + target : str + {"cpu", "gpu", "cuda"} + auto_broadcast : bool + auto_broadcast=True allows one to implement broadcast computation + without considering whether dimension size equals to one. + TVM maps buffer[i][j][k] -> buffer[i][0][k] if dimension i's shape equals 1. + Returns + ------- + fdef : function + A wrapped operator definition function, which returns (schedule, [tensors]) + """ + assert name is not None and len(name) > 0 + target = "cpu" if target is None else target + def _defop(func): + opdef = OpDef(func, name, target, auto_broadcast, **kwargs) + __OP_DEF__.append(opdef) + return opdef + return _defop + diff --git a/contrib/tvmop/prepare_tvm.sh b/contrib/tvmop/prepare_tvm.sh new file mode 100644 index 000000000000..7ebe256446a4 --- /dev/null +++ b/contrib/tvmop/prepare_tvm.sh @@ -0,0 +1,63 @@ +# 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. + +#!/bin/sh + +LLVM_VERSION="8.0.0" +LLVM_ROOT="http://releases.llvm.org/${LLVM_VERSION}/" +LLVM_PKG="clang+llvm-${LLVM_VERSION}-x86_64-linux-gnu" + +os=`uname` +if [ "$os" = "Linux" ] && [ "$(arch)" = "x86_64" ]; then + DISTRIB_ID=$(cat /etc/*-release | grep DISTRIB_ID | sed 's/DISTRIB_ID=//g' | tr '[:upper:]' '[:lower:]') + DISTRIB_RELEASE=$(cat /etc/*-release | grep DISTRIB_RELEASE | sed 's/DISTRIB_RELEASE=//g' | tr '[:upper:]' '[:lower:]') + if [ "$DISTRIB_ID" = "ubuntu" ]; then + LLVM_PKG=${LLVM_PKG}-${DISTRIB_ID}-${DISTRIB_RELEASE} + else + echo "Downloading LLVM only supports Ubuntu. Please download manually." + exit 1 + fi +else + echo "Cannot identify operating system. Try downloading LLVM manually." + exit 1 +fi + +LLVM_URL=${LLVM_ROOT}${LLVM_PKG}.tar.xz + +TVM_PATH=`dirname $0`/../../3rdparty/tvm +DST=${TVM_PATH}/build +rm -rf $DST +mkdir -p $DST +DST=`cd $DST; pwd` + +if [ -x "$(command -v curl)" ]; then + curl -L -o "${DST}/${LLVM_PKG}.tar.xz" "$LLVM_URL" +elif [ -x "$(command -v wget)" ]; then + wget -O "${DST}/${LLVM_PKG}.tar.xz" "$LLVM_URL" +else + echo "curl or wget not available" + exit 1 +fi + +if [ \! $? ]; then + echo "Download from $LLVM_URL to $DST failed" + exit 1 +fi + +tar -xf "$DST/${LLVM_PKG}.tar.xz" -C $DST +mv $DST/${LLVM_PKG} $DST/llvm +echo "Downloaded and unpacked LLVM libraries to $DST" diff --git a/contrib/tvmop/utils.py b/contrib/tvmop/utils.py new file mode 100644 index 000000000000..0b2416b4f3ae --- /dev/null +++ b/contrib/tvmop/utils.py @@ -0,0 +1,20 @@ +# 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 +AllTypes = ["float32", "float64", "float16", "uint8", "int8", "int32", "int64"] +RealTypes = ["float32", "float64", "float16"] diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h index bd30e44f910c..058f859ae7bc 100644 --- a/include/mxnet/c_api.h +++ b/include/mxnet/c_api.h @@ -506,6 +506,15 @@ MXNET_DLL int MXGetGPUMemoryInformation64(int dev, uint64_t *free_mem, uint64_t */ MXNET_DLL int MXGetVersion(int *out); +/*! + * \brief Load TVM operator from the binary library + * \param libpath TVM operators lib file + * \return 0 when success, -1 when failure happens + */ +#if MXNET_USE_TVM_OP +MXNET_DLL int MXLoadTVMOp(const char *libpath); +#endif // MXNET_USE_TVM_OP + //------------------------------------- // Part 1: NDArray creation and deletion diff --git a/include/mxnet/libinfo.h b/include/mxnet/libinfo.h index 8b58a398c673..1972688c7739 100644 --- a/include/mxnet/libinfo.h +++ b/include/mxnet/libinfo.h @@ -127,6 +127,10 @@ #define MXNET_USE_INT64_TENSOR_SIZE MSHADOW_INT64_TENSOR_SIZE #endif +#ifndef MXNET_USE_TVM_OP +#define MXNET_USE_TVM_OP 0 +#endif + namespace mxnet { namespace features { // Check compile flags such as CMakeLists.txt @@ -185,6 +189,9 @@ enum : unsigned { SIGNAL_HANDLER, DEBUG, + // TVM operator + TVM_OP, + // size indicator MAX_FEATURES }; diff --git a/make/config.mk b/make/config.mk index 4bddb8ba461d..982d15b19656 100644 --- a/make/config.mk +++ b/make/config.mk @@ -62,6 +62,9 @@ ADD_LDFLAGS = # the additional compile flags you want to add ADD_CFLAGS = +# whether to build operators written in TVM +USE_TVM_OP = 0 + #--------------------------------------------- # matrix computation libraries for CPU/GPU #--------------------------------------------- diff --git a/make/osx.mk b/make/osx.mk index 0b5842e59524..25f3ba6df55b 100644 --- a/make/osx.mk +++ b/make/osx.mk @@ -53,6 +53,9 @@ ADD_LDFLAGS = # the additional compile flags you want to add ADD_CFLAGS = +# whether to build operators written in TVM +USE_TVM_OP = 0 + #--------------------------------------------- # matrix computation libraries for CPU/GPU #--------------------------------------------- diff --git a/python/mxnet/base.py b/python/mxnet/base.py index 73fae4876873..bf8026359d02 100644 --- a/python/mxnet/base.py +++ b/python/mxnet/base.py @@ -16,7 +16,7 @@ # under the License. # coding: utf-8 -# pylint: disable=invalid-name, no-member, trailing-comma-tuple, bad-mcs-classmethod-argument, unnecessary-pass +# pylint: disable=invalid-name, no-member, trailing-comma-tuple, bad-mcs-classmethod-argument, unnecessary-pass, wrong-import-position """ctypes library of mxnet and helper functions.""" from __future__ import absolute_import @@ -734,3 +734,8 @@ def write_all_str(module_file, module_all_list): ctypes.pythonapi.PyCapsule_New.restype = ctypes.py_object ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p + +from .runtime import Features +if Features().is_enabled("TVM_OP"): + _LIB_TVM_OP = libinfo.find_lib_path("libtvmop") + check_call(_LIB.MXLoadTVMOp(c_str(_LIB_TVM_OP[0]))) diff --git a/python/mxnet/libinfo.py b/python/mxnet/libinfo.py index ff795f914a4b..fb0859b9954a 100644 --- a/python/mxnet/libinfo.py +++ b/python/mxnet/libinfo.py @@ -23,7 +23,7 @@ import logging -def find_lib_path(): +def find_lib_path(prefix='libmxnet'): """Find MXNet dynamic library files. Returns @@ -61,13 +61,13 @@ def find_lib_path(): dll_path[0:0] = [p.strip() for p in os.environ['LD_LIBRARY_PATH'].split(":")] if os.name == 'nt': os.environ['PATH'] = os.path.dirname(__file__) + ';' + os.environ['PATH'] - dll_path = [os.path.join(p, 'libmxnet.dll') for p in dll_path] + dll_path = [os.path.join(p, prefix + '.dll') for p in dll_path] elif platform.system() == 'Darwin': - dll_path = [os.path.join(p, 'libmxnet.dylib') for p in dll_path] + \ - [os.path.join(p, 'libmxnet.so') for p in dll_path] + dll_path = [os.path.join(p, prefix + '.dylib') for p in dll_path] + \ + [os.path.join(p, prefix + '.so') for p in dll_path] else: dll_path.append('../../../') - dll_path = [os.path.join(p, 'libmxnet.so') for p in dll_path] + dll_path = [os.path.join(p, prefix + '.so') for p in dll_path] lib_path = [p for p in dll_path if os.path.exists(p) and os.path.isfile(p)] if len(lib_path) == 0: raise RuntimeError('Cannot find the MXNet library.\n' + diff --git a/src/c_api/c_api.cc b/src/c_api/c_api.cc index 35bd3eeb477a..5207bdfa444b 100644 --- a/src/c_api/c_api.cc +++ b/src/c_api/c_api.cc @@ -48,6 +48,7 @@ #include "./c_api_common.h" #include "../operator/custom/custom-inl.h" #include "../operator/tensor/matrix_op-inl.h" +#include "../operator/tvmop/op_module.h" #include "../common/utils.h" using namespace mxnet; @@ -159,6 +160,14 @@ int MXGetVersion(int *out) { API_END(); } +#if MXNET_USE_TVM_OP +int MXLoadTVMOp(const char *libpath) { + API_BEGIN(); + tvm::runtime::TVMOpModule::Get()->Load(libpath); + API_END(); +} +#endif // MXNET_USE_TVM_OP + int MXNDArrayCreateNone(NDArrayHandle *out) { API_BEGIN(); *out = new NDArray(); diff --git a/src/libinfo.cc b/src/libinfo.cc index f67b45ed1c14..b31d7e4301e0 100644 --- a/src/libinfo.cc +++ b/src/libinfo.cc @@ -89,6 +89,9 @@ class FeatureSet { feature_bits.set(INT64_TENSOR_SIZE, MXNET_USE_INT64_TENSOR_SIZE); feature_bits.set(SIGNAL_HANDLER, MXNET_USE_SIGNAL_HANDLER); + // TVM operators + feature_bits.set(TVM_OP, MXNET_USE_TVM_OP); + #ifndef NDEBUG feature_bits.set(DEBUG); #endif @@ -159,6 +162,7 @@ const std::vector EnumNames::names = { "INT64_TENSOR_SIZE", "SIGNAL_HANDLER", "DEBUG", + "TVM_OP", }; } // namespace features diff --git a/src/operator/contrib/tvmop/ufunc.cc b/src/operator/contrib/tvmop/ufunc.cc new file mode 100644 index 000000000000..faba671322e2 --- /dev/null +++ b/src/operator/contrib/tvmop/ufunc.cc @@ -0,0 +1,66 @@ +/* + * 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. + */ + +/*! + * Copyright (c) 2019 by Contributors + * \file ufunc.cc + * \brief + * \author Yizhi Liu + */ +#ifdef MXNET_USE_TVM_OP +#include +#include +#include +#include +#include "../../tensor/elemwise_binary_broadcast_op.h" +#include "../../tvmop/op_module.h" +#include "../../tensor/elemwise_binary_op.h" + +namespace mxnet { +namespace op { + +static constexpr char func_vadd_cpu[] = "vadd"; +static constexpr char func_vadd_gpu[] = "cuda_vadd"; + +template +void TVMBroadcastCompute(const nnvm::NodeAttrs& attrs, + const mxnet::OpContext& ctx, + const std::vector& inputs, + const std::vector& req, + const std::vector& outputs) { + CHECK_EQ(inputs.size(), 2U); + CHECK_EQ(outputs.size(), 1U); + tvm::runtime::TVMOpModule::Get()->Call(func, ctx, {inputs[0], inputs[1], outputs[0]}); +} + +NNVM_REGISTER_OP(_contrib_tvm_vadd) + .set_num_inputs(2) + .set_num_outputs(1) + .add_argument("a", "NDArray-or-Symbol", "first input") + .add_argument("b", "NDArray-or-Symbol", "second input") + .set_attr("FInferShape", BinaryBroadcastShape) + .set_attr("FInferType", mxnet::op::ElemwiseType<2, 1>) + .set_attr("FCompute", mxnet::op::TVMBroadcastCompute) +#if MXNET_USE_CUDA + .set_attr("FCompute", mxnet::op::TVMBroadcastCompute); +#endif // MXNET_USE_CUDA + +} // namespace op +} // namespace mxnet +#endif // MXNET_USE_TVM_OP diff --git a/src/operator/tvmop/op_module.cc b/src/operator/tvmop/op_module.cc new file mode 100644 index 000000000000..d1d1c1d45d53 --- /dev/null +++ b/src/operator/tvmop/op_module.cc @@ -0,0 +1,117 @@ +/* + * 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. + */ + +/*! + * Copyright (c) 2019 by Contributors + * \file op_module.cc + * \brief Invoke registered TVM operators. + * \author Yizhi Liu + */ +#if MXNET_USE_TVM_OP +#include +#include +#include +#include +#include +#include "op_module.h" + +using namespace tvm::runtime; + +namespace tvm { +namespace runtime { + +void TVMOpModule::Load(const std::string &filepath) { + static const PackedFunc *f_load = Registry::Get("module._LoadFromFile"); + std::lock_guard lock(mutex_); + Module module = (*f_load)(filepath, ""); + module_ptr_ = std::make_shared(); + *module_ptr_ = module; +} + +PackedFunc GetFunction(const std::shared_ptr &module, + const std::string &op_name, + const std::vector &args) { + std::ostringstream func_name; + func_name << op_name; + for (const auto &arg : args) { + switch (arg.type_flag_) { + case mshadow::kFloat32: + func_name << "float32"; + break; + case mshadow::kFloat64: + func_name << "float64"; + break; + case mshadow::kFloat16: + func_name << "float16"; + break; + case mshadow::kUint8: + func_name << "uint8"; + break; + case mshadow::kInt32: + func_name << "int32"; + break; + case mshadow::kInt8: + func_name << "int8"; + break; + case mshadow::kInt64: + func_name << "int64"; + break; + default: + LOG(FATAL) << "Unknown dtype " << arg.type_flag_; + } + func_name << "_" << arg.shape_.ndim(); + } + return module->GetFunction(func_name.str(), false); +} + +void TVMOpModule::Call(const std::string &func_name, + const mxnet::OpContext& ctx, + const std::vector &args) { + std::vector type_codes; + std::vector values; + + type_codes.resize(args.size()); + values.resize(args.size()); + for (size_t i = 0; i < args.size(); ++i) { + type_codes[i] = kArrayHandle; + values[i].v_handle = const_cast(&(args[i].dltensor())); + } + + TVMArgs tvm_args(&values[0], &type_codes[0], args.size()); + TVMRetValue rv; + +#if MXNET_USE_CUDA + int dev_type = (ctx.run_ctx.ctx.dev_type == mxnet::Context::DeviceType::kGPU) ? kDLGPU : kDLCPU; + int dev_id = ctx.run_ctx.ctx.dev_id; + if (dev_type == kDLGPU) { + void *stream = static_cast(ctx.run_ctx.get_stream()->stream_); + TVMSetStream(dev_type, dev_id, stream); + } +#endif + GetFunction(module_ptr_, func_name, args).CallPacked(tvm_args, &rv); +#if MXNET_USE_CUDA + if (dev_type == kDLGPU) { + TVMSetStream(dev_type, dev_id, nullptr); + } +#endif +} + +} // namespace runtime +} // namespace tvm +#endif // MXNET_USE_TVM_OP diff --git a/src/operator/tvmop/op_module.h b/src/operator/tvmop/op_module.h new file mode 100644 index 000000000000..04e97ef51f4e --- /dev/null +++ b/src/operator/tvmop/op_module.h @@ -0,0 +1,63 @@ +/* + * 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. + */ + +/*! + * Copyright (c) 2019 by Contributors + * \file op_module.h + * \brief Invoke registered TVM operators. + * \author Yizhi Liu + */ +#ifndef MXNET_OPERATOR_TVMOP_OP_MODULE_H_ +#define MXNET_OPERATOR_TVMOP_OP_MODULE_H_ + +#if MXNET_USE_TVM_OP +#include +#include +#include +#include +#include + +namespace tvm { +namespace runtime { + +class Module; +class TVMOpModule { + public: + // Load TVM operators binary + void Load(const std::string& filepath); + + void Call(const std::string& func_name, + const mxnet::OpContext& ctx, + const std::vector& args); + + static TVMOpModule *Get() { + static TVMOpModule inst; + return &inst; + } + + private: + std::mutex mutex_; + std::shared_ptr module_ptr_; +}; + +} // namespace runtime +} // namespace tvm + +#endif // MXNET_USE_TVM_OP +#endif // MXNET_OPERATOR_TVMOP_OP_MODULE_H_ diff --git a/tests/python/gpu/test_operator_gpu.py b/tests/python/gpu/test_operator_gpu.py index 5b4f81d96065..dcd6acce8ca5 100644 --- a/tests/python/gpu/test_operator_gpu.py +++ b/tests/python/gpu/test_operator_gpu.py @@ -44,6 +44,7 @@ from test_ndarray import * from test_subgraph_op import * from test_contrib_operator import test_multibox_target_op +from test_tvm_op import * set_default_context(mx.gpu(0)) del test_support_vector_machine_l1_svm # noqa diff --git a/tests/python/unittest/test_tvm_op.py b/tests/python/unittest/test_tvm_op.py new file mode 100644 index 000000000000..3ab2a25bc20e --- /dev/null +++ b/tests/python/unittest/test_tvm_op.py @@ -0,0 +1,38 @@ +# 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. + +import mxnet as mx +from mxnet.test_utils import same, rand_shape_nd +from mxnet.runtime import Features +from common import with_seed + +_features = Features() + +@with_seed() +def test_tvm_broadcast_add(): + if _features.is_enabled("TVM_OP"): + a_shape = rand_shape_nd(4) + b_shape = (1,) + a_shape[1:2] + (1, 1) + a = mx.nd.normal(shape=a_shape) + b = mx.nd.normal(shape=b_shape) + c = mx.nd.contrib.tvm_vadd(a, b) + c_np = a.asnumpy() + b.asnumpy() + assert same(c.asnumpy(), c_np) + +if __name__ == '__main__': + import nose + nose.runmodule()