Added scalar type support

Added SetScalar api to support scalar input Ovxlib need update to support gather with scalar index Type: New Feature Signed-off-by: Feiyue Chen <[email protected]>
VeriSilicon · Oct 26, 2023 · d4187b3 · d4187b3
1 parent 1008179
commit d4187b3
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git a/include/tim/vx/tensor.h b/include/tim/vx/tensor.h
@@ -161,7 +161,9 @@ class Tensor {
   virtual void unmap() = 0;
   virtual bool IsPlaceHolder() = 0;
   virtual bool IsConstTensor() = 0;
+  virtual bool IsScalar() = 0;
   virtual bool SaveTensorToTextByFp32(std::string filename) = 0;
+  virtual void SetScalar(int8_t is_scalar) = 0;
   virtual void* ConvertTensorToData(uint8_t* tensorData) = 0;
   virtual float* ConvertTensorToFloat32Data() = 0;
 };

diff --git a/src/tim/vx/tensor.cc b/src/tim/vx/tensor.cc
@@ -213,6 +213,14 @@ float* TensorImpl::ConvertTensorToFloat32Data() {
       graph_->graph(), vsi_nn_GetTensor(graph_->graph(), id_));
 }
 
+void TensorImpl::SetScalar(int8_t is_scalar) {
+  bool retn = vsi_nn_SetTensorIsScalar(vsi_nn_GetTensor(graph_->graph(), id_),is_scalar);
+  if (retn != VSI_SUCCESS) {
+    VSILOGE("Setting scalar fail!");
+  }
+  return;
+}
+
 bool TensorImpl::SwapHandle(void* new_ptr, bool is_new_ptr_malloc_by_ovxlib,
                             void** old_ptr) {
   bool retn = true;

diff --git a/src/tim/vx/tensor_private.h b/src/tim/vx/tensor_private.h
@@ -60,9 +60,14 @@ class TensorImpl : public Tensor {
   bool IsConstTensor() override {
     return spec_.attr_ == tim::vx::TensorAttribute::CONSTANT;
   }
+  bool IsScalar() override {
+    return vsi_nn_GetTensorIsScalar(vsi_nn_GetTensor(graph_->graph(), id_));
+  }
   bool SaveTensorToTextByFp32(std::string filename) override;
   void* ConvertTensorToData(uint8_t* tensorData) override;
   float* ConvertTensorToFloat32Data() override;
+  void SetScalar(int8_t is_scalar) override;
+
   GraphImpl* graph_;
   vsi_nn_tensor_id_t id_;
   TensorSpec spec_;
@@ -114,6 +119,9 @@ class TensorPlaceholder : public Tensor {
   bool IsConstTensor() override {
     return spec_.attr_ == tim::vx::TensorAttribute::CONSTANT;
   }
+  bool IsScalar() override {
+    return false;
+  }
   bool SaveTensorToTextByFp32(std::string filename) override {
     (void)filename;
     return false;
@@ -124,6 +132,7 @@ class TensorPlaceholder : public Tensor {
   }
   float* ConvertTensorToFloat32Data() override { return nullptr; }
 
+  void SetScalar(int8_t is_scalar) override { (void) is_scalar; return; }
   vsi_nn_tensor_id_t id_;
   TensorSpec spec_;
 };