[RUNTIME] Implement TVMDSOOp(TensorFlow custom op) for TVM runtime

CMakeLists.txt

@@ -42,6 +42,7 @@ tvm_option(USE_MSVC_MT "Build with MT" OFF)
 tvm_option(USE_MICRO "Build with Micro" OFF)
 tvm_option(INSTALL_DEV "Install compiler infrastructure" OFF)
 tvm_option(HIDE_PRIVATE_SYMBOLS "Compile with -fvisibility=hidden." OFF)
+tvm_option(USE_TFOP "Build with TensorFlow TVMDSOOp" OFF)
 
 # 3rdparty libraries
 tvm_option(DLPACK_PATH "Path to DLPACK" "3rdparty/dlpack/include")
@@ -260,6 +261,7 @@ include(cmake/modules/contrib/Sort.cmake)
 include(cmake/modules/contrib/NNPack.cmake)
 include(cmake/modules/contrib/HybridDump.cmake)
 include(cmake/modules/contrib/TFLite.cmake)
+include(cmake/modules/contrib/TFOP.cmake)
 
 if(NOT MSVC)
   include(CheckCXXCompilerFlag)

cmake/config.cmake

@@ -193,3 +193,6 @@ set(USE_VTA_FPGA OFF)
 
 # Whether to build the example external runtime module
 set(USE_EXAMPLE_EXT_RUNTIME OFF)
+
+# Whether to build the TensorFlow TVMDSOOp module
+set(USE_TFOP OFF)

cmake/modules/contrib/TFOP.cmake

@@ -0,0 +1,46 @@
+
+if(NOT USE_TFOP STREQUAL "OFF")
+
+  # If want build this directly comment out below lines.
+  # if ("${TVM_HOME}" STREQUAL "")
+  #   message(FATAL_ERROR "TVM_HOME is not defined")
+  # else()
+  #  message("Use TVM_HOME=\"${TVM_HOME}\"")
+  #endif()
+  # include_directories(${TVM_HOME}/include)
+  # include_directories(${TVM_HOME}/3rdparty/dlpack/include)
+  # include_directories(${TVM_HOME}/3rdparty/dmlc-core/include)
+  # set(TFTVM_LINK_FLAGS  -ltvm_runtime -L${TVM_HOME}/build)
+
+  execute_process(COMMAND python -c "import tensorflow as tf; print(' '.join(tf.sysconfig.get_compile_flags()))"
+    OUTPUT_VARIABLE TF_COMPILE_FLAGS_STR
+    RESULT_VARIABLE TF_STATUS)
+  if (NOT ${TF_STATUS} EQUAL 0)
+    message(FATAL_ERROR "Fail to get TensorFlow compile flags")
+  endif()
+
+  execute_process(COMMAND python -c "import tensorflow as tf; print(' '.join(tf.sysconfig.get_link_flags()))"
+    OUTPUT_VARIABLE TF_LINK_FLAGS_STR
+    RESULT_VARIABLE TF_STATUS)
+  if (NOT ${TF_STATUS} EQUAL 0)
+    message(FATAL_ERROR "Fail to get TensorFlow link flags")
+  endif()
+
+  string(REGEX REPLACE "\n" " " TF_FLAGS "${TF_COMPILE_FLAGS} ${TF_LINK_FLAGS}")
+  separate_arguments(TF_COMPILE_FLAGS UNIX_COMMAND ${TF_COMPILE_FLAGS_STR})
+  separate_arguments(TF_LINK_FLAGS UNIX_COMMAND ${TF_LINK_FLAGS_STR})
+
+
+  set(OP_LIBRARY_NAME tvm_dso_op)
+  file(GLOB_RECURSE TFTVM_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/src/contrib/tf_op/*.cc)
+  add_library(${OP_LIBRARY_NAME} SHARED ${TFTVM_SRCS})
+  set_target_properties(${OP_LIBRARY_NAME} PROPERTIES PREFIX "")
+  set(TFTVM_LINK_FLAGS  -ltvm -L${CMAKE_CURRENT_BINARY_DIR})
+  add_dependencies(${OP_LIBRARY_NAME} tvm) 
+
+  set(TFTVM_COMPILE_FLAGS  -O2 -ldl -g)
+  target_compile_options(${OP_LIBRARY_NAME} PUBLIC ${TFTVM_COMPILE_FLAGS} ${TF_COMPILE_FLAGS})
+  target_link_options(${OP_LIBRARY_NAME} PUBLIC ${TFTVM_LINK_FLAGS} ${TF_LINK_FLAGS})
+
+endif()
+

python/tvm/contrib/tf_op/__init__.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.
+
+from . import module
+
+Module = module.Module

python/tvm/contrib/tf_op/module.py

@@ -0,0 +1,104 @@
+# 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 tensorflow as tf
+from tensorflow.python.framework import load_library
+
+
+class Module():
+
+  def __init__(self, lib_path):
+    self.lib_path = lib_path
+
+  def func(self, name, output_dtype=None, output_shape=None):
+    return Func(self.lib_path, name, output_dtype, output_shape)
+
+  def __getitem__(self, func_name):
+    return self.func(func_name)
+
+
+class Func():
+
+  def __init__(self, lib_path, func_name, output_dtype, output_shape):
+    self.lib_path = lib_path
+    self.func_name = func_name
+    self.output_dtype = output_dtype
+
+    # const(0) indicate invalid dynamic shape
+    self.dynamic_output_shape = tf.constant(0, tf.int64)
+    self.static_output_shape = None
+    self.has_static_output_shape = False  # extra flag is required
+
+    if self._is_static_shape(output_shape):
+      self.static_output_shape = output_shape
+      self.has_static_output_shape = True
+    elif output_shape is not None:
+      self.dynamic_output_shape = self._pack_shape_tensor(output_shape)
+
+    # TODO: support non-xpu device 
+    #self.device = device
+    # delay initialization to called first time, where num input arguments is known
+    self.tvm_dso_op = None
+    self.module = load_library.load_op_library('tvm_dso_op.so')
+
+  def apply(self, *params):
+    if self.tvm_dso_op is None:
+      num_inputs = len(params)
+      self.tvm_dso_op = getattr(self.module, "tvm_dso_op%s" % num_inputs)
+
+    return self.tvm_dso_op(*params, 
+                           dynamic_output_shape=self.dynamic_output_shape,
+                           static_output_shape=self.static_output_shape,
+                           has_static_output_shape=self.has_static_output_shape, 
+                           lib_path=self.lib_path, 
+                           func_name=self.func_name, 
+                           output_dtype=self.output_dtype)
+
+  def __call__(self, *params):
+    return self.apply(*params)
+
+  def _is_static_shape(self, shape):
+    if shape is None or not isinstance(shape, list):
+      return False
+    for d in shape:
+      if not isinstance(d, int):
+        return False
+      if d < 0:
+        raise Exception("Negative dimension is illegal: %d" % d)
+    return True
+
+  def _pack_shape_tensor(self, shape):
+    if isinstance(shape, tf.Tensor):
+      if shape.dtype == tf.int32:
+        shape = tf.cast(shape, tf.int64)
+      return shape
+    elif isinstance(shape, list):
+      shape_dims = []
+      for d in shape:
+        if isinstance(d, int):
+          shape_dims.append(tf.constant(d, tf.int64))
+        elif isinstance(d, tf.Tensor) and len(d.shape) == 0:
+          if d.dtype == tf.int32:
+            d = tf.cast(d, tf.int64)
+          shape_dims.append(d)
+        else:
+          raise TypeError("Input shape dimension is neither scala tensor nor int")
+      return tf.stack(shape_dims) 
+    else:
+      raise TypeError("Input shape is neither tensor nor list")
+
+
+

src/contrib/tf_op/index_seq.h

@@ -0,0 +1,44 @@
+/**
+ * Refer to std::index_sequence (since c++14)
+ * Utilities to invoke variadic function with template <size_t N> 
+ */
+#ifndef __TFTVM_INDEX_SEQ__
+#define __TFTVM_INDEX_SEQ__
+
+template <std::size_t ...>
+struct IndexSeq {};
+
+template <std::size_t N, std::size_t ... Tail>
+struct IndexSeqHelper : public IndexSeqHelper<N-1U, N-1U, Tail...> {};
+
+template <std::size_t ... Tail>
+struct IndexSeqHelper<0U, Tail ...> {
+    using type = IndexSeq<Tail ...>;
+};
+
+template <std::size_t N>
+using make_index_sequence = typename IndexSeqHelper<N>::type;
+
+
+template <typename F, typename T, std::size_t N, std::size_t... Idx>
+void apply_variadic_impl(F f, T(&t)[N], IndexSeq<Idx...>) {
+    f(t[Idx]...);
+}
+
+template <typename F, typename T, std::size_t N>
+void apply_variadic(F f, T(&t)[N]) {
+    apply_variadic_impl(f, t, make_index_sequence<N>{});
+}
+
+template <typename F, typename T, std::size_t N, std::size_t... Idx>
+void apply_variadic_by_ptrs_impl(F f, T(&t)[N], IndexSeq<Idx...>) {
+    f(&t[Idx]...);
+}
+
+template <typename F, typename T, std::size_t N>
+void apply_variadic_by_ptrs(F f, T(&t)[N]) {
+    apply_variadic_by_ptrs_impl(f, t, make_index_sequence<N>{});
+}
+
+#endif
+