From da11b53751d251fef76e4bdd4264f80c66d08d6c Mon Sep 17 00:00:00 2001
From: zhaozhixuan <zhaozhixuan2@hisilicon.com>
Date: Wed, 12 May 2021 14:25:32 +0800
Subject: [PATCH] Set storage shape to single_op executor.

---
 ge/hybrid/executor/hybrid_model_executor.cc |  2 +-
 ge/single_op/single_op.cc                   | 22 ++++++++++++++++++++++
 tests/ut/ge/single_op/single_op_unittest.cc |  7 +++++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/ge/hybrid/executor/hybrid_model_executor.cc b/ge/hybrid/executor/hybrid_model_executor.cc
index 2ab4ed5d..60db0d33 100755
--- a/ge/hybrid/executor/hybrid_model_executor.cc
+++ b/ge/hybrid/executor/hybrid_model_executor.cc
@@ -50,7 +50,7 @@ Status HybridModelExecutor::Execute(HybridModelExecutor::ExecuteArgs &args) {
   auto root_graph_item = model_->GetRootGraphItem();
   GE_CHECK_NOTNULL(root_graph_item);
 
-  if (root_graph_item->IsDynamic()) {
+  if (root_graph_item->IsDynamic() && !model_->IsSingleOp()) {
     GE_CHK_STATUS_RET(CheckInputShapeByShapeRange(root_graph_item, args),
                       "[%s] check input node shape by shape range failed.",
                       root_graph_item->GetName().c_str());
diff --git a/ge/single_op/single_op.cc b/ge/single_op/single_op.cc
index 4b3f17cf..b9cd4e85 100755
--- a/ge/single_op/single_op.cc
+++ b/ge/single_op/single_op.cc
@@ -113,6 +113,27 @@ Status UpdateInputsBufferAddr(StreamResource *stream_resource, rtStream_t stream
   return SUCCESS;
 }
 
+Status ModifyTensorDesc(GeTensorDesc &tensor) {
+  int64_t storage_format_val = static_cast<Format>(FORMAT_RESERVED);
+  (void)AttrUtils::GetInt(tensor, ge::ATTR_NAME_STORAGE_FORMAT, storage_format_val);
+  auto storage_format = static_cast<Format>(storage_format_val);
+  if (storage_format != FORMAT_RESERVED) {
+    std::vector<int64_t> storage_shape;
+    if (!AttrUtils::GetListInt(tensor, ge::ATTR_NAME_STORAGE_SHAPE, storage_shape)) {
+      GELOGE(ACL_ERROR_GE_INTERNAL_ERROR, "[Get][storage_shape]failed while storage_format was set.");
+      REPORT_INNER_ERROR("E19999", "Get storage_shape failed while storage_format was set.");
+      return ACL_ERROR_GE_INTERNAL_ERROR;
+    }
+
+    GELOGD("Storage format set. update shape to [%s], and original shape to [%s]",
+           GeShape(storage_shape).ToString().c_str(), tensor.GetShape().ToString().c_str());
+    tensor.SetShape(GeShape(std::move(storage_shape)));
+    tensor.SetFormat(std::move(storage_format));
+  }
+
+  return SUCCESS;
+}
+
 Status InitHybridModelArgs(const std::vector<DataBuffer> &input_buffers,
                            const std::vector<DataBuffer> &output_buffers,
                            const std::vector<GeTensorDesc> &inputs_desc,
@@ -126,6 +147,7 @@ Status InitHybridModelArgs(const std::vector<DataBuffer> &input_buffers,
   for (auto &tensor_desc : inputs_desc) {
     auto desc = MakeShared<GeTensorDesc>(tensor_desc);
     GE_CHECK_NOTNULL(desc);
+    GE_CHK_STATUS_RET_NOLOG(ModifyTensorDesc(*desc));
     args.input_desc.emplace_back(desc);
   }
   return SUCCESS;
diff --git a/tests/ut/ge/single_op/single_op_unittest.cc b/tests/ut/ge/single_op/single_op_unittest.cc
index 8c2f6e51..bdd6401d 100644
--- a/tests/ut/ge/single_op/single_op_unittest.cc
+++ b/tests/ut/ge/single_op/single_op_unittest.cc
@@ -159,5 +159,12 @@ TEST_F(UtestSingleOp, test_singleop_execute_async2) {
   single_op.hybrid_model_executor_.reset(new (std::nothrow)hybrid::HybridModelExecutor(single_op.hybrid_model_.get(), 0, stream));
   EXPECT_EQ(single_op.running_param_->mem_base, nullptr);
   EXPECT_EQ(single_op.tasks_.size(), 0);
+
+  GeTensorDesc tensor;
+  int64_t storage_format_val = static_cast<Format>(FORMAT_NCHW);
+  AttrUtils::SetInt(tensor, "storage_format", storage_format_val);
+  std::vector<int64_t> storage_shape{1, 1, 1, 1};
+  AttrUtils::SetListInt(tensor, "storage_shape", storage_shape);
+  single_op.inputs_desc_.emplace_back(tensor);
   EXPECT_EQ(single_op.ExecuteAsync(input_buffers, output_buffers), PARAM_INVALID);
 }
\ No newline at end of file