From ab0929662b4e17cc95aa909f09cb82b9c3f71ae3 Mon Sep 17 00:00:00 2001
From: heliqi <1101791222@qq.com>
Date: Wed, 28 Dec 2022 21:58:32 +0800
Subject: [PATCH] [Runtime]FDTensor pybind add from_dlpack interface (#1001)

add from_dlpack
---
 fastdeploy/pybind/fd_tensor.cc | 69 ++++++++++++++++++++++++++++++++++
 fastdeploy/pybind/runtime.cc   |  3 +-
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/fastdeploy/pybind/fd_tensor.cc b/fastdeploy/pybind/fd_tensor.cc
index bf24dec6b..0f1d145b3 100644
--- a/fastdeploy/pybind/fd_tensor.cc
+++ b/fastdeploy/pybind/fd_tensor.cc
@@ -166,6 +166,74 @@ pybind11::capsule FDTensorToDLPack(FDTensor& fd_tensor) {
       static_cast<void*>(dlpack_tensor), "dltensor", &DeleteUnusedDltensor);
 }
 
+FDTensor FDTensorFromDLPack(const std::string& name,
+                            const pybind11::capsule& dlpack_tensor) {
+  DLManagedTensor* dl_managed_tensor =
+      static_cast<DLManagedTensor*>(dlpack_tensor.get_pointer());
+
+  void* memory_ptr = dl_managed_tensor->dl_tensor.data;
+  memory_ptr = reinterpret_cast<char*>(memory_ptr) +
+               dl_managed_tensor->dl_tensor.byte_offset;
+
+  int64_t* strides = dl_managed_tensor->dl_tensor.strides;
+
+  int ndim = dl_managed_tensor->dl_tensor.ndim;
+  std::vector<int64_t> dims(
+      dl_managed_tensor->dl_tensor.shape,
+      dl_managed_tensor->dl_tensor.shape + ndim);
+
+  // Check if the input is contiguous and in C order
+  if (strides != nullptr) {
+    int64_t calculated_stride{1};
+    bool is_contiguous_c_order = true;
+    for (size_t i = 1; i < dims.size(); i++) {
+      if (strides[ndim - i] != calculated_stride) {
+        is_contiguous_c_order = false;
+        break;
+      }
+
+      calculated_stride *= dims[ndim - i];
+    }
+
+    FDASSERT(is_contiguous_c_order,
+        "DLPack tensor is not contiguous. Only contiguous DLPack "
+        "tensors that are stored in C-Order are supported.");
+  }
+
+  Device device;
+  int32_t device_id = -1;
+  bool is_pinned_memory = false;
+  switch (dl_managed_tensor->dl_tensor.device.device_type) {
+    case DLDeviceType::kDLCUDA:
+      device = Device::GPU;
+      device_id = dl_managed_tensor->dl_tensor.device.device_id;
+      break;
+    case DLDeviceType::kDLCPU:
+      device = Device::CPU;
+      break;
+    case DLDeviceType::kDLCUDAHost:
+      device = Device::CPU;
+      is_pinned_memory = true;
+      break;
+    default:
+      FDASSERT(false,
+          ("DLDevice type " +
+          std::to_string(dl_managed_tensor->dl_tensor.device.device_type) +
+          " is not support by Python backend.").c_str());
+      break;
+  }
+
+  FDDataType dtype =
+      DlpackToFDType(dl_managed_tensor->dl_tensor.dtype);
+
+  PyCapsule_SetName(dlpack_tensor.ptr(), "used_dlpack");
+  FDTensor fd_tensor(name);
+  fd_tensor.SetExternalData(
+    dims, dtype, memory_ptr, device, device_id
+  );
+  fd_tensor.is_pinned_memory = is_pinned_memory;
+  return fd_tensor;
+}
 
 void BindFDTensor(pybind11::module& m) {
   pybind11::class_<FDTensor>(m, "FDTensor")
@@ -182,6 +250,7 @@ void BindFDTensor(pybind11::module& m) {
         PyArrayToTensor(pyarray, &self, share_buffer);
       })
       .def("to_dlpack", &FDTensorToDLPack)
+      .def("from_dlpack",&FDTensorFromDLPack)
       .def("print_info", &FDTensor::PrintInfo);
 }
 
diff --git a/fastdeploy/pybind/runtime.cc b/fastdeploy/pybind/runtime.cc
index 5c299da0e..cbbb7f2e8 100755
--- a/fastdeploy/pybind/runtime.cc
+++ b/fastdeploy/pybind/runtime.cc
@@ -214,7 +214,8 @@ void BindRuntime(pybind11::module& m) {
       .def("infer",
            [](Runtime& self, std::vector<FDTensor>& inputs) {
              std::vector<FDTensor> outputs;
-             return self.Infer(inputs, &outputs);
+             self.Infer(inputs, &outputs);
+             return outputs;
            })
       .def("bind_input_tensor", &Runtime::BindInputTensor)
       .def("infer", [](Runtime& self) { self.Infer(); })