Skip to content

Commit ed1998b

Browse files
yzhliuRohit Kumar Srivastava
yzhliu
authored and
Rohit Kumar Srivastava
committed
Infra to use tvm write op kernels (apache#15550)
* intra to use tvm write op kernels * add cmake support for tvm op * fix header lint * cleanup USE_TVM_OP logic in Makefile * add doc, cmake def, etc. * fix doc * test rand shape * add with_seed to test * improve err msg. add #if
1 parent 0a8fa0b commit ed1998b

File tree

26 files changed

+874
-9
lines changed

26 files changed

+874
-9
lines changed

‎3rdparty/dlpack

‎3rdparty/dmlc-core

‎3rdparty/tvm

Submodule tvm updated from 21935dc to afd4b3e

‎CMakeLists.txt

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ mxnet_option(USE_MXNET_LIB_NAMING "Use MXNet library naming conventions." ON)
4343
mxnet_option(USE_GPROF "Compile with gprof (profiling) flag" OFF)
4444
mxnet_option(USE_CXX14_IF_AVAILABLE "Build with C++14 if the compiler supports it" OFF)
4545
mxnet_option(USE_VTUNE "Enable use of Intel Amplifier XE (VTune)" OFF) # one could set VTUNE_ROOT for search path
46+
mxnet_option(USE_TVM_OP "Enable use of TVM operator build system." OFF)
4647
mxnet_option(ENABLE_CUDA_RTC "Build with CUDA runtime compilation support" ON)
4748
mxnet_option(BUILD_CPP_EXAMPLES "Build cpp examples" ON)
4849
mxnet_option(INSTALL_EXAMPLES "Install the example source files." OFF)
@@ -734,6 +735,28 @@ if(USE_DIST_KVSTORE)
734735
list(APPEND mxnet_LINKER_LIBS ${pslite_LINKER_LIBS})
735736
endif()
736737

738+
if(USE_TVM_OP)
739+
add_definitions(-DMXNET_USE_TVM_OP=1)
740+
list(APPEND mxnet_LINKER_LIBS ${CMAKE_CURRENT_BINARY_DIR}/3rdparty/tvm/libtvm_runtime.so)
741+
include(cmake/BuildTVM.cmake)
742+
add_subdirectory("3rdparty/tvm")
743+
744+
if(NOT Python3_EXECUTABLE)
745+
find_package(PythonInterp 3 REQUIRED)
746+
set(Python3_EXECUTABLE ${PYTHON_EXECUTABLE} CACHE FILEPATH "Path to the python3 executable")
747+
if(NOT Python3_EXECUTABLE)
748+
message(FATAL_ERROR "No python3 interpreter found to build TVM operators")
749+
endif()
750+
endif()
751+
752+
add_custom_command(TARGET mxnet POST_BUILD
753+
COMMAND ${CMAKE_COMMAND} -E env
754+
PYTHONPATH="${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm/python:${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm/topi/python:${CMAKE_CURRENT_SOURCE_DIR}/contrib"
755+
LD_LIBRARY_PATH="${CMAKE_CURRENT_BINARY_DIR}/3rdparty/tvm/build"
756+
${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/contrib/tvmop/compile.py -o${CMAKE_CURRENT_BINARY_DIR}/libtvmop.so
757+
)
758+
endif()
759+
737760
target_link_libraries(mxnet PUBLIC ${mxnet_LINKER_LIBS})
738761

739762
if(USE_PLUGINS_WARPCTC)

‎Makefile

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,14 @@ ifndef AMALGAMATION_PATH
5252
AMALGAMATION_PATH = $(ROOTDIR)/amalgamation
5353
endif
5454

55+
ifndef TVM_PATH
56+
TVM_PATH = $(TPARTYDIR)/tvm
57+
endif
58+
59+
ifndef LLVM_PATH
60+
LLVM_PATH = $(TVM_PATH)/build/llvm
61+
endif
62+
5563
ifneq ($(USE_OPENMP), 1)
5664
export NO_OPENMP = 1
5765
endif
@@ -589,6 +597,35 @@ $(DMLC_CORE)/libdmlc.a: DMLCCORE
589597
DMLCCORE:
590598
+ cd $(DMLC_CORE); $(MAKE) libdmlc.a USE_SSE=$(USE_SSE) config=$(ROOTDIR)/$(config); cd $(ROOTDIR)
591599

600+
ifeq ($(USE_TVM_OP), 1)
601+
LIB_DEP += lib/libtvm_runtime.so lib/libtvmop.so
602+
CFLAGS += -I$(TVM_PATH)/include -DMXNET_USE_TVM_OP=1
603+
LDFLAGS += -L$(TVM_PATH)/build -ltvm_runtime
604+
605+
TVM_USE_CUDA := OFF
606+
ifeq ($(USE_CUDA), 1)
607+
TVM_USE_CUDA := ON
608+
ifneq ($(USE_CUDA_PATH), NONE)
609+
TVM_USE_CUDA := $(USE_CUDA_PATH)
610+
endif
611+
endif
612+
lib/libtvm_runtime.so:
613+
echo "Compile TVM"
614+
[ -e $(LLVM_PATH)/bin/llvm-config ] || sh $(ROOTDIR)/contrib/tvmop/prepare_tvm.sh; \
615+
cd $(TVM_PATH)/build; \
616+
cmake -DUSE_LLVM="$(LLVM_PATH)/bin/llvm-config" \
617+
-DUSE_SORT=OFF -DUSE_CUDA=$(TVM_USE_CUDA) -DUSE_CUDNN=OFF ..; \
618+
$(MAKE) VERBOSE=1; \
619+
cp $(TVM_PATH)/build/libtvm_runtime.so $(ROOTDIR)/lib/libtvm_runtime.so; \
620+
cd $(ROOTDIR)
621+
622+
lib/libtvmop.so: lib/libtvm_runtime.so $(wildcard contrib/tvmop/*/*.py contrib/tvmop/*.py)
623+
echo "Compile TVM operators"
624+
PYTHONPATH=$(TVM_PATH)/python:$(TVM_PATH)/topi/python:$(ROOTDIR)/contrib:$PYTHONPATH \
625+
LD_LIBRARY_PATH=lib \
626+
python3 $(ROOTDIR)/contrib/tvmop/compile.py -o $(ROOTDIR)/lib/libtvmop.so
627+
endif
628+
592629
NNVM_INC = $(wildcard $(NNVM_PATH)/include/*/*.h)
593630
NNVM_SRC = $(wildcard $(NNVM_PATH)/src/*/*/*.cc $(NNVM_PATH)/src/*/*.cc $(NNVM_PATH)/src/*.cc)
594631
$(NNVM_PATH)/lib/libnnvm.a: $(NNVM_INC) $(NNVM_SRC)
@@ -726,6 +763,7 @@ clean: rclean cyclean $(EXTRA_PACKAGES_CLEAN)
726763
cd $(DMLC_CORE); $(MAKE) clean; cd -
727764
cd $(PS_PATH); $(MAKE) clean; cd -
728765
cd $(NNVM_PATH); $(MAKE) clean; cd -
766+
cd $(TVM_PATH); $(MAKE) clean; cd -
729767
cd $(AMALGAMATION_PATH); $(MAKE) clean; cd -
730768
$(RM) -r $(patsubst %, %/*.d, $(EXTRA_OPERATORS)) $(patsubst %, %/*/*.d, $(EXTRA_OPERATORS))
731769
$(RM) -r $(patsubst %, %/*.o, $(EXTRA_OPERATORS)) $(patsubst %, %/*/*.o, $(EXTRA_OPERATORS))

‎cmake/BuildTVM.cmake

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
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+
18+
message(STATUS "Prepare external packages for TVM...")
19+
execute_process(COMMAND "${CMAKE_CURRENT_SOURCE_DIR}/contrib/tvmop/prepare_tvm.sh")
20+
21+
# Whether enable ROCM runtime
22+
#
23+
# Possible values:
24+
# - ON: enable ROCM with cmake's auto search
25+
# - OFF: disable ROCM
26+
# - /path/to/rocm: use specific path to rocm
27+
set(USE_ROCM OFF)
28+
29+
# Whether enable SDAccel runtime
30+
set(USE_SDACCEL OFF)
31+
32+
# Whether enable Intel FPGA SDK for OpenCL (AOCL) runtime
33+
set(USE_AOCL OFF)
34+
35+
# Whether enable OpenCL runtime
36+
set(USE_OPENCL OFF)
37+
38+
# Whether enable Metal runtime
39+
set(USE_METAL OFF)
40+
41+
# Whether enable Vulkan runtime
42+
#
43+
# Possible values:
44+
# - ON: enable Vulkan with cmake's auto search
45+
# - OFF: disable vulkan
46+
# - /path/to/vulkan-sdk: use specific path to vulkan-sdk
47+
set(USE_VULKAN OFF)
48+
49+
# Whether enable OpenGL runtime
50+
set(USE_OPENGL OFF)
51+
52+
# Whether to enable SGX runtime
53+
#
54+
# Possible values for USE_SGX:
55+
# - /path/to/sgxsdk: path to Intel SGX SDK
56+
# - OFF: disable SGX
57+
#
58+
# SGX_MODE := HW|SIM
59+
set(USE_SGX OFF)
60+
set(SGX_MODE "SIM")
61+
set(RUST_SGX_SDK "/path/to/rust-sgx-sdk")
62+
63+
# Whether enable RPC runtime
64+
set(USE_RPC ON)
65+
66+
# Whether embed stackvm into the runtime
67+
set(USE_STACKVM_RUNTIME OFF)
68+
69+
# Whether enable tiny embedded graph runtime.
70+
set(USE_GRAPH_RUNTIME ON)
71+
72+
# Whether enable additional graph debug functions
73+
set(USE_GRAPH_RUNTIME_DEBUG OFF)
74+
75+
# Whether build with LLVM support
76+
# Requires LLVM version >= 4.0
77+
#
78+
# Possible values:
79+
# - ON: enable llvm with cmake's find search
80+
# - OFF: disable llvm
81+
# - /path/to/llvm-config: enable specific LLVM when multiple llvm-dev is available.
82+
set(USE_LLVM "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm/build/llvm/bin/llvm-config")
83+
84+
#---------------------------------------------
85+
# Contrib libraries
86+
#---------------------------------------------
87+
# Whether use BLAS, choices: openblas, mkl, atlas, apple
88+
set(USE_BLAS none)
89+
90+
# /path/to/mkl: mkl root path when use mkl blas library
91+
# set(USE_MKL_PATH /opt/intel/mkl) for UNIX
92+
# set(USE_MKL_PATH ../IntelSWTools/compilers_and_libraries_2018/windows/mkl) for WIN32
93+
set(USE_MKL_PATH none)
94+
95+
# Whether use contrib.random in runtime
96+
set(USE_RANDOM OFF)
97+
98+
# Whether use NNPack
99+
set(USE_NNPACK OFF)
100+
101+
# Whether use CuDNN
102+
if(USE_CUDNN AND USE_CUDA)
103+
detect_cuDNN()
104+
if(HAVE_CUDNN)
105+
set(USE_CUDNN ON)
106+
else()
107+
set(USE_CUDNN OFF)
108+
endif()
109+
else()
110+
set(USE_CUDNN OFF)
111+
endif()
112+
113+
# Whether use cuBLAS
114+
set(USE_CUBLAS OFF)
115+
116+
# Whether use MIOpen
117+
set(USE_MIOPEN OFF)
118+
119+
# Whether use MPS
120+
set(USE_MPS OFF)
121+
122+
# Whether use rocBlas
123+
set(USE_ROCBLAS OFF)
124+
125+
# Whether use contrib sort
126+
set(USE_SORT OFF)
127+
128+
# Build ANTLR parser for Relay text format
129+
set(USE_ANTLR OFF)
130+
131+
# Build TSIM for VTA
132+
set(USE_VTA_TSIM OFF)
133+
134+
# Whether use Relay debug mode
135+
set(USE_RELAY_DEBUG OFF)

‎contrib/tvmop/__init__.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
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+
18+
# coding: utf-8
19+
from .opdef import defop
20+
from .utils import AllTypes, RealTypes
21+
22+
from . import basic

‎contrib/tvmop/basic/__init__.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
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+
18+
# coding: utf-8
19+
from . import ufunc

‎contrib/tvmop/basic/ufunc.py

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
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+
18+
# coding: utf-8
19+
import tvm
20+
from .. import defop, AllTypes
21+
22+
def compute_add(dtype, ndim):
23+
A = tvm.placeholder([tvm.var() for _ in range(ndim)], name='A', dtype=dtype)
24+
B = tvm.placeholder([tvm.var() for _ in range(ndim)], name='B', dtype=dtype)
25+
C = tvm.compute([tvm.var() for _ in range(ndim)],
26+
lambda *index: A[index] + B[index], name='C')
27+
s = tvm.create_schedule(C.op)
28+
return s, A, B, C
29+
30+
@defop(name="vadd", target="cpu", auto_broadcast=True,
31+
dtype=AllTypes, ndim=list(range(1, 6)))
32+
def vadd(dtype, ndim):
33+
s, A, B, C = compute_add(dtype, ndim)
34+
axes = [axis for axis in C.op.axis]
35+
fused = s[C].fuse(*axes)
36+
s[C].parallel(fused)
37+
38+
return s, [A, B, C]
39+
40+
@defop(name="cuda_vadd", target="cuda", auto_broadcast=True,
41+
dtype=["float32", "float64"], ndim=list(range(1, 6)))
42+
def vadd_gpu(dtype, ndim):
43+
s, A, B, C = compute_add(dtype, ndim)
44+
s = tvm.create_schedule(C.op)
45+
axes = [axis for axis in C.op.axis]
46+
fused = s[C].fuse(*axes)
47+
bx, tx = s[C].split(fused, factor=64)
48+
s[C].bind(bx, tvm.thread_axis("blockIdx.x"))
49+
s[C].bind(tx, tvm.thread_axis("threadIdx.x"))
50+
return s, [A, B, C]

‎contrib/tvmop/compile.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
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+
18+
# coding: utf-8
19+
"""TVM Operator compile entry point"""
20+
import tvm
21+
22+
import os
23+
import argparse
24+
from tvmop.opdef import __OP_DEF__
25+
26+
def get_target(device):
27+
if device == "cpu":
28+
return "llvm"
29+
elif device == "cuda" or device == "gpu":
30+
return "cuda"
31+
assert False, "Unknown device " + device
32+
33+
34+
if __name__ == "__main__":
35+
import sys
36+
sys.path.append(os.path.dirname(sys.path[0]))
37+
parser = argparse.ArgumentParser(description="Generate tvm operators")
38+
parser.add_argument("-o", action="store", required=True, dest="target_path",
39+
help="Target path which stores compiled library")
40+
arguments = parser.parse_args()
41+
42+
func_list_llvm = []
43+
func_list_cuda = []
44+
45+
# TODO: attach instruction features to the library, e.g., avx-512, etc.
46+
for operator_def in __OP_DEF__:
47+
for sch, args in operator_def.invoke_all():
48+
if tvm.module.enabled(get_target(operator_def.target)):
49+
func_list = func_list_llvm if operator_def.target == "cpu" else func_list_cuda
50+
func_lower = tvm.lower(sch, args,
51+
name=operator_def.get_op_name(args),
52+
binds=operator_def.get_binds(args))
53+
func_list.append(func_lower)
54+
55+
lowered_funcs = {get_target("cpu") : func_list_llvm}
56+
if len(func_list_cuda) > 0:
57+
lowered_funcs[get_target("cuda")] = func_list_cuda
58+
func_binary = tvm.build(lowered_funcs, name="tvmop")
59+
func_binary.export_library(arguments.target_path)

0 commit comments

Comments
 (0)