From 8b1ec74f07d76153be50502dbad625fc078cbe56 Mon Sep 17 00:00:00 2001
From: chxin66 <57057788+chxin66@users.noreply.github.com>
Date: Sun, 21 Nov 2021 22:46:20 +0800
Subject: [PATCH] support DMAbuffer (#214)

Signed-off-by: Chen Xin <jack.chen@verisilicon.com>
---
 include/tim/vx/graph.h      |  7 ++++---
 include/tim/vx/tensor.h     |  4 ++++
 src/tim/vx/graph.cc         |  5 +++++
 src/tim/vx/graph_private.h  |  2 ++
 src/tim/vx/tensor.cc        | 15 +++++++++++++--
 src/tim/vx/tensor_private.h |  2 ++
 6 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/include/tim/vx/graph.h b/include/tim/vx/graph.h
index 4fb66ab..4527320 100644
--- a/include/tim/vx/graph.h
+++ b/include/tim/vx/graph.h
@@ -26,13 +26,12 @@
 
 #include <memory>
 #include <vector>
-
 namespace tim {
 namespace vx {
 
 class Tensor;
 struct TensorSpec;
-
+struct DmaBufferDesc;
 class Operation;
 
 class Graph {
@@ -43,6 +42,8 @@ class Graph {
   virtual std::shared_ptr<Tensor> CreateTensor(const TensorSpec& spec,
                                                const void* data = nullptr) = 0;
 
+  virtual std::shared_ptr<Tensor> CreateTensor(const TensorSpec& spec,
+                                               const DmaBufferDesc& dmafd) = 0;
   /// Create a placeholder tensor for optional inputs of operations
   virtual std::shared_ptr<Tensor> CreateTensorPlaceHolder() = 0;
 
@@ -70,7 +71,7 @@ class Graph {
 
   virtual const std::vector<std::shared_ptr<Operation>> GetConsumersOp(
       std::shared_ptr<Tensor> tensor) const = 0;
-  
+
   virtual void PrintGraph() const = 0;
 
  protected:
diff --git a/include/tim/vx/tensor.h b/include/tim/vx/tensor.h
index d9f8812..7297fcc 100644
--- a/include/tim/vx/tensor.h
+++ b/include/tim/vx/tensor.h
@@ -140,6 +140,10 @@ struct TensorSpec {
   Quantization quantization_;
 };
 
+struct DmaBufferDesc {
+  int64_t fd;
+};
+
 class Tensor {
  public:
   virtual ~Tensor() {}
diff --git a/src/tim/vx/graph.cc b/src/tim/vx/graph.cc
index e8db240..abda269 100644
--- a/src/tim/vx/graph.cc
+++ b/src/tim/vx/graph.cc
@@ -105,6 +105,11 @@ std::shared_ptr<Tensor> GraphImpl::CreateTensor(const TensorSpec& spec,
   return std::make_shared<TensorImpl>(this, spec, data);
 }
 
+std::shared_ptr<Tensor> GraphImpl::CreateTensor(const TensorSpec& spec,
+                                                const DmaBufferDesc& dmafd) {
+  return std::make_shared<TensorImpl>(this, spec, dmafd);
+}
+
 std::shared_ptr<Tensor> GraphImpl::CreateTensorPlaceHolder() {
   if (!tensor_placeholder_) {
     tensor_placeholder_ = std::make_shared<TensorPlaceholder>(this);
diff --git a/src/tim/vx/graph_private.h b/src/tim/vx/graph_private.h
index 426f2fd..241536f 100644
--- a/src/tim/vx/graph_private.h
+++ b/src/tim/vx/graph_private.h
@@ -62,6 +62,8 @@ class GraphImpl : public Graph {
   /// Implement parents' virtual functions
    std::shared_ptr<Tensor> CreateTensor(const TensorSpec& spec,
                                        const void* data = nullptr) override;
+   std::shared_ptr<Tensor> CreateTensor(const TensorSpec& spec,
+                                               const DmaBufferDesc& dmafd) override;
    std::shared_ptr<Tensor> CreateTensorPlaceHolder() override;
     bool Compile() override;
 
diff --git a/src/tim/vx/tensor.cc b/src/tim/vx/tensor.cc
index d12ed79..1124b6e 100644
--- a/src/tim/vx/tensor.cc
+++ b/src/tim/vx/tensor.cc
@@ -85,6 +85,14 @@ TensorImpl::TensorImpl(Graph* graph, const TensorSpec& spec, const void* data)
   Init();
 }
 
+TensorImpl::TensorImpl(Graph* graph, const TensorSpec& spec, const DmaBufferDesc& dmafd)
+    : graph_(reinterpret_cast<GraphImpl*>(graph)),
+      id_(VSI_NN_TENSOR_ID_NA),
+      spec_(spec),
+      fd_(dmafd.fd) {
+  Init();
+}
+
 TensorImpl::~TensorImpl() {}
 
 bool TensorImpl::CopyDataToTensor(const void* data, uint32_t size_in_bytes) {
@@ -183,8 +191,11 @@ bool TensorImpl::Init() {
 
   if ((spec_.attr_ & TensorAttribute::INPUT) ||
       (spec_.attr_ & TensorAttribute::OUTPUT)) {
-    id_ = vsi_nn_AddTensorFromHandle(graph_->graph(), VSI_NN_TENSOR_ID_AUTO,
-                                     &attr, nullptr);
+    id_ = vsi_nn_AddTensorFromHandle(graph_->graph(), VSI_NN_TENSOR_ID_AUTO,     // DMABUF's fd is created by TensorFromHandle as input or output,
+                                     &attr, fd_ != -1 ? (uint8_t*)fd_ : nullptr);// and cannot be set to const
+    if (fd_ != -1) {
+      attr.vsi_memory_type = VSI_MEMORY_TYPE_DMABUF;
+    }
   } else {
     id_ = vsi_nn_AddTensor(graph_->graph(), VSI_NN_TENSOR_ID_AUTO, &attr,
                            nullptr);
diff --git a/src/tim/vx/tensor_private.h b/src/tim/vx/tensor_private.h
index 367586e..ce26d84 100644
--- a/src/tim/vx/tensor_private.h
+++ b/src/tim/vx/tensor_private.h
@@ -33,6 +33,7 @@ namespace vx {
 class TensorImpl : public Tensor {
  public:
   TensorImpl(Graph* graph, const TensorSpec& spec, const void* data = nullptr);
+  TensorImpl(Graph* graph, const TensorSpec& spec, const DmaBufferDesc& dmafd);
   ~TensorImpl();
 
   bool Init();
@@ -56,6 +57,7 @@ class TensorImpl : public Tensor {
   vsi_nn_tensor_id_t id_;
   TensorSpec spec_;
   const void* data_;
+  int64_t fd_{-1};
 };
 
 class TensorPlaceholder : public Tensor {