tvm-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From tqc...@apache.org
Subject [incubator-tvm] branch master updated: Load platform specific lib for tvmdsoop instead of only so (#5542)
Date Fri, 08 May 2020 19:42:34 GMT
This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-tvm.git


The following commit(s) were added to refs/heads/master by this push:
     new 15a4218  Load platform specific lib for tvmdsoop instead of only so (#5542)
15a4218 is described below

commit 15a421880f252b18d3636606234d8e20b91047f5
Author: tobe <tobeg3oogle@gmail.com>
AuthorDate: Sat May 9 03:42:22 2020 +0800

    Load platform specific lib for tvmdsoop instead of only so (#5542)
---
 python/tvm/contrib/tf_op/module.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/python/tvm/contrib/tf_op/module.py b/python/tvm/contrib/tf_op/module.py
index f13670e..fd0ed0c 100644
--- a/python/tvm/contrib/tf_op/module.py
+++ b/python/tvm/contrib/tf_op/module.py
@@ -17,6 +17,7 @@
 """Module container of TensorFlow TVMDSO op"""
 import tensorflow as tf
 from tensorflow.python.framework import load_library
+from tensorflow.python import platform
 
 
 class OpModule:
@@ -67,7 +68,7 @@ class TensorFunc:
         elif output_shape is not None:
             self.dynamic_output_shape = self._pack_shape_tensor(output_shape)
 
-        self.module = load_library.load_op_library('tvm_dso_op.so')
+        self.module = self._load_platform_specific_library("tvm_dso_op")
         self.tvm_dso_op = self.module.tvm_dso_op
 
     def apply(self, *params):
@@ -82,6 +83,16 @@ class TensorFunc:
     def __call__(self, *params):
         return self.apply(*params)
 
+    def _load_platform_specific_library(self, lib_name):
+        system = platform.system()
+        if system == "Darwin":
+            lib_file_name = lib_name + ".dylib"
+        elif system == "Windows":
+            lib_file_name = lib_name + ".dll"
+        else:
+            lib_file_name = lib_name + ".so"
+        return load_library.load_op_library(lib_file_name)
+
     def _is_static_shape(self, shape):
         if shape is None or not isinstance(shape, list):
             return False


Mime
View raw message