update build instrucion

Author: jermainewang

CMakeLists.txt

@@ -4,7 +4,8 @@ project(tfdlpack C CXX)
 
 add_definitions(-std=c++11 -fPIC)
 include(cmake/util/FindCUDA.cmake)	
-option(NO_CUDA "Only build TFDLPACK with cpu" OFF)
+
+option(USE_CUDA "Build TF-DLPACK with CUDA support" ON)
 
 if(NOT PYTHON_EXECUTABLE)
   execute_process(
@@ -35,12 +36,15 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O2 ${TF_CFLAGS}")
 set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -O2 -D_GLIBCXX_USE_CXX11_ABI=0 ${TF_CFLAGS}")
 set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -O0 -g3 -ggdb -D_GLIBCXX_USE_CXX11_ABI=0 ${TF_CFLAGS}")
 
-if (NOT NO_CUDA)
+if (USE_CUDA)
+  message(STATUS "Build with CUDA.")
   add_definitions(-DTFDLPACK_USE_CUDA)
   find_cuda(ON REQUIRED)
   if(NOT CUDA_FOUND)	
     message(FATAL_ERROR "Cannot find CUDA.")	
   endif()
+else()
+  message(STATUS "Build without CUDA.")
 endif()
 
 include_directories(third_party/dlpack/include)
@@ -49,10 +53,10 @@ file(GLOB SRC
   src/*.cc
 )
 
-if (NO_CUDA)
-  add_library(tfdlpack SHARED ${SRC})
-else()
+if (USE_CUDA)
   cuda_add_library(tfdlpack SHARED ${SRC})
+else()
+  add_library(tfdlpack SHARED ${SRC})
 endif()
 
 target_link_libraries(tfdlpack ${TF_LFLAGS})

README.md

@@ -36,16 +36,18 @@ Build
 ```
 mkdir build
 cd build
-cmake ..
+cmake ..  # To build without CUDA, add -DUSE_CUDA=OFF
 make -j4
 ```
 
-so file path is now fixed in `python/tfdlpack/__init__.py`
-Need to change manually
+Export the library path:
+```bash
+export TF_DLPACK_LIBRARY_PATH=/path/to/tf-dlpack/repo/build
+```
 
-And export the python path to `import tfdlpack`
+Export python path to `import tfdlpack`
 ```bash
-export PYTHONPATH=/home/ubuntu/dev/tfdlpack/python/:${PYTHONPATH}
+export PYTHONPATH=/path/to/tf-dlpack/repo/python/:${PYTHONPATH}
 ```
 
 ## License

python/tfdlpack/core.py

@@ -10,7 +10,6 @@
 # version number
 __version__ = libinfo.__version__
 
-# print(libinfo.find_lib_path()[0])
 dlpack_ops = load_library.load_op_library(libinfo.find_lib_path()[0])
 _to_dlpack_address = dlpack_ops.to_dlpack
 _from_dlpack = dlpack_ops.from_dlpack

src/from_dlpack_kernel.cc

@@ -25,16 +25,14 @@ inline bool IsAligned(size_t alignment, void *data_ptr) {
 
 class DLPackAllocator : public Allocator {
  public:
-  static constexpr size_t kAllocatorAlignment = 64;
-
   explicit DLPackAllocator(DLManagedTensor *dlm_tensor) {
     dlm_tensor_ = dlm_tensor;
     data_ = dlm_tensor->dl_tensor.data;
 
     // Shape
     shape_ = TensorShape();
-    int ndim = dlm_tensor->dl_tensor.ndim;
-    int64_t *shape = dlm_tensor->dl_tensor.shape;
+    const int ndim = dlm_tensor->dl_tensor.ndim;
+    const int64_t *shape = dlm_tensor->dl_tensor.shape;
     for (int i = 0; i < ndim; i++) {
       shape_.AddDim(shape[i]);
     }
@@ -61,7 +59,6 @@ class DLPackAllocator : public Allocator {
   void DeallocateRaw(void *ptr) {
     // This would lead to double free, haven't figure out the problem
     dlm_tensor_->deleter(const_cast<DLManagedTensor *>(dlm_tensor_));
-    // std::cout << "Deconstruct dlpack tensor" << std::endl;
     delete this;
   }
 
@@ -78,6 +75,8 @@ class DLPackAllocator : public Allocator {
   int64 num_elements_;
   TensorShape shape_;
   Status allocation_status_;
+
+  TF_DISALLOW_COPY_AND_ASSIGN(DLPackAllocator);
 };
 
 class FromDLPackOP : public OpKernel {

tests/test_zero_copy.py

@@ -2,10 +2,10 @@
 import torch as th
 from torch.utils.dlpack import from_dlpack, to_dlpack
 from tfdlpack import from_dlpack as tf_from_dlpack
-import gpustat
 import pytest
 
 def get_gpu_memory_used():
+    import gpustat
     gpu_query = gpustat.GPUStatCollection.new_query()
     gmem_used = gpu_query[0].memory_used
     return gmem_used
@@ -28,4 +28,4 @@ def test_zero_copy():
     print(m3)
 
 if __name__ == "__main__":
-    test_zero_copy()
+    test_zero_copy()