feat: (1) Support zero copy from tf tensor to dlpack (2) Fix pymodul…

CMakeLists.txt

@@ -0,0 +1,46 @@
+cmake_minimum_required(VERSION 3.2)
+project(tftvm)
+set(CMAKE_VERBOSE_MAKEFILE ON)
+
+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)
+
+
+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}")
+message("Use TensorFlow flags=\"${TF_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 src/*.cc)
+add_library(${OP_LIBRARY_NAME} SHARED ${TFTVM_SRCS})
+set_target_properties(${OP_LIBRARY_NAME} PROPERTIES PREFIX "")
+
+set(TFTVM_COMPILE_FLAGS  -O2 -ldl -g)
+set(TFTVM_LINK_FLAGS  -ltvm_runtime -L${TVM_HOME}/build)
+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})
+

build.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+mkdir -p build
+cd build
+cmake .. -DTVM_HOME=${TVM_HOME}
+make

python/tf_op/module.py

@@ -41,7 +41,7 @@ def __init__(self, lib_path, func_name, output_dtype, output_shape, device):
     self.tvm_dso_op  = tvm_dso_op.tvm_dso_op
 
   def apply(self, *params):
-    return self.tvm_dso_op(params, lib_path=self.lib_path, func_name=self.func_name, output_dtype=self.output_dtype, output_shape=self.output_shape, device=self.device)
+    return self.tvm_dso_op(*params, lib_path=self.lib_path, func_name=self.func_name, output_dtype=self.output_dtype, output_shape=self.output_shape, device=self.device)
 
   def __call__(self, *params):
-    return self.apply(params)
+    return self.apply(*params)