Add support for DLPack Tensors (triton-inference-server#62)

* Move triton_pb.Tensor to C++

* Add support for DLPack GPU tensor

* Add support for conversion between dlpack datatypes and triton

* Add support for output tensors in GPU

* Fix bug when the output tensors contain one of the input tensors

* Add clean up for DLPack tensors

* Fix deadlock when init takes long

* Fix DLPack for CPU tensors

* Clean up

* Fix up

* Add FORCE_CPU_ONLY_INPUT_TENSORS

* Minor fixes for TYPE_STRING and contiguous tensors

* Review edits

* Manually adjust the offset for cudaIpcMemHandle

* Print appropriate error when the tensor is not contiguous

* Refactor communication between the main process and the stub process (triton-inference-server#65)

* Fix CPU_PINNED_MEMORY

* Delete shared memory region if there is an exception

* Minor improvements

* Fix shutdown and identity tensor return

* Add more comment

Author: Tabrizian

CMakeLists.txt

@@ -84,6 +84,17 @@ FetchContent_Declare(
 )
 FetchContent_MakeAvailable(pybind11)
 
+#
+# DLPack
+#
+FetchContent_Declare(
+  dlpack
+  GIT_REPOSITORY "https://github.com/dmlc/dlpack"
+  GIT_TAG "v0.5"
+  GIT_SHALLOW ON
+)
+FetchContent_MakeAvailable(dlpack)
+
 #
 # Boost
 #
@@ -99,6 +110,15 @@ ExternalProject_Add(
 )
 set(boostorg_INCLUDE_DIRS "${CMAKE_BINARY_DIR}/boost/")
 
+#
+# CUDA
+#
+if(${TRITON_ENABLE_GPU})
+  find_package(CUDAToolkit REQUIRED)
+  message(STATUS "Using CUDA ${CUDA_VERSION}")
+  set(CUDA_NVCC_FLAGS -std=c++11)
+endif() # TRITON_ENABLE_GPU
+
 find_package(ZLIB REQUIRED)
 find_package(Threads REQUIRED)
 
@@ -118,9 +138,13 @@ add_library(
 
 add_executable(
   triton-python-backend-stub
+  src/pb_stub_utils.cc
+  src/pb_stub_utils.h
   src/pb_stub.cc
   src/pb_utils.cc
   src/pb_utils.h
+  src/pb_tensor.cc
+  src/pb_tensor.h
   src/shm_manager.cc
   src/shm_manager.h
 )
@@ -162,11 +186,26 @@ target_link_libraries(
   PRIVATE
    Threads::Threads
    pybind11::embed
+   dlpack
    triton-backend-utils    # from repo-backend
    -larchive               # libarchive
    -lrt                    # shared memory 
 )
 
+if(${TRITON_ENABLE_GPU})
+  target_link_libraries(
+    triton-python-backend-stub
+    PUBLIC
+      CUDA::cuda_driver
+  )
+
+  target_link_libraries(
+    triton-python-backend
+    PUBLIC
+      CUDA::cuda_driver
+  )
+endif() # TRITON_ENABLE_GPU
+
 set_target_properties(
   triton-python-backend PROPERTIES
   POSITION_INDEPENDENT_CODE ON