Fix dump for known-shaped subgraph

ge/hybrid/executor/hybrid_model_executor.cc

@@ -33,9 +33,6 @@ HybridModelExecutor::~HybridModelExecutor() {
   if (context_.rt_gen_context != nullptr) {
     (void) rtCtxDestroy(context_.rt_gen_context);
   }
-  if (context_.global_step != nullptr) {
-    (void) rtFree(context_.global_step);
-  }
 }
 
 Status HybridModelExecutor::Init() {
@@ -49,7 +46,6 @@ Status HybridModelExecutor::Execute(HybridModelExecutor::ExecuteArgs &args) {
   GELOGD("Start to execute model.");
   auto root_graph_item = model_->GetRootGraphItem();
   GE_CHECK_NOTNULL(root_graph_item);
-
   GE_CHK_RT_RET(rtMemcpyAsync(context_.global_step, sizeof(uint64_t), &context_.iteration,
                               sizeof(uint64_t), RT_MEMCPY_HOST_TO_DEVICE_EX, context_.stream));
   SubgraphExecutor executor(model_->GetRootGraphItem(), &context_);
@@ -102,8 +98,8 @@ Status HybridModelExecutor::InitExecutionContext() {
   GE_CHK_RT_RET(rtCtxGetCurrent(&context_.rt_context));
   GE_CHK_RT_RET(rtCtxCreate(&context_.rt_gen_context, RT_CTX_GEN_MODE, 0));
   GE_CHK_RT_RET(rtCtxSetCurrent(context_.rt_context));
-  GE_CHK_RT_RET(rtMalloc(&context_.global_step, sizeof(uint64_t), RT_MEMORY_HBM));
 
+  context_.global_step = model_->GetGlobalStep();
   context_.stream = stream_;
   context_.model = model_;
   context_.is_eos_ = false;

ge/hybrid/model/hybrid_model.cc

@@ -357,5 +357,12 @@ TensorValue *HybridModel::GetTensor(const NodePtr &node) const {
 
   return GetVariable(node->GetName());
 }
+
+void *HybridModel::GetGlobalStep() const {
+  if (global_step_ == nullptr) {
+    return nullptr;
+  }
+  return global_step_->GetData();
+}
 }  // namespace hybrid
 }  // namespace ge

ge/hybrid/model/hybrid_model.h

@@ -45,6 +45,8 @@ class HybridModel {
     return root_runtime_param_.session_id;
   }
 
+  void *GetGlobalStep() const;
+
   GeModelPtr GetGeModel(const NodePtr &node) const;
 
   NodeItem *MutableNodeItem(const NodePtr &node);
@@ -157,6 +159,7 @@ class HybridModel {
   std::map<string, std::unique_ptr<TensorBuffer>> weight_buffer_map_;
   RuntimeParam root_runtime_param_;
   string om_name_;
+  std::unique_ptr<TensorBuffer> global_step_;
 };
 }  // namespace hybrid
 }  // namespace ge

ge/hybrid/model/hybrid_model_builder.cc

@@ -130,8 +130,8 @@ HybridModelBuilder::HybridModelBuilder(HybridModel &hybrid_model)
 
 Status HybridModelBuilder::Build() {
   GE_CHK_STATUS_RET(ValidateParams(), "Failed to validate GeRootModel");
-  hybrid_model_.model_name_ = ge_root_model_->GetRootGraph()->GetName();
   GELOGI("[%s] Start to build hybrid model.", GetGraphName());
+  GE_CHK_STATUS_RET(InitHybridModel(), "Failed to int model");
   GE_CHK_STATUS_RET(InitRuntimeParams(), "[%s] Failed to InitRuntimeParams", GetGraphName());
   GE_CHK_STATUS_RET(RecoverGraphUnknownFlag(), "[%s] Failed to RecoverGraphUnknownFlag", GetGraphName());
   GE_CHK_STATUS_RET(IndexSpecialNodes(), "[%s] Failed to index nodes", GetGraphName());
@@ -151,7 +151,7 @@ Status HybridModelBuilder::Build() {
 
 Status HybridModelBuilder::BuildForSingleOp() {
   GE_CHK_STATUS_RET(ValidateParams(), "Failed to validate GeRootModel");
-  hybrid_model_.model_name_ = ge_root_model_->GetRootGraph()->GetName();
+  GE_CHK_STATUS_RET(InitHybridModel(), "Failed to int model");
   GELOGI("[%s] Start to build hybrid model.", GetGraphName());
   auto ret = ge_root_model_->GetSubgraphInstanceNameToModel();
   const GeModelPtr ge_model = ret[ge_root_model_->GetRootGraph()->GetName()];
@@ -2113,5 +2113,14 @@ Status HybridModelBuilder::ParseDependentByParallelGroup() {
   }
   return SUCCESS;
 }
+
+Status HybridModelBuilder::InitHybridModel() {
+  hybrid_model_.model_name_ = ge_root_model_->GetRootGraph()->GetName();
+  auto allocator = NpuMemoryAllocator::GetAllocator();
+  GE_CHECK_NOTNULL(allocator);
+  hybrid_model_.global_step_ = TensorBuffer::Create(allocator, sizeof(int64_t));
+  GE_CHECK_NOTNULL(hybrid_model_.global_step_);
+  return SUCCESS;
+}
 }  // namespace hybrid
 }  // namespace ge

ge/hybrid/model/hybrid_model_builder.h

@@ -80,6 +80,8 @@ class HybridModelBuilder {
   Status AssignUninitializedConstantOps();
   Status InitConstantOps();
   Status InitVariableTensors();
+  Status InitHybridModel();
+
   Status LoadDynamicSubgraph(ComputeGraph &graph, bool is_root_graph);
   Status ParseVarOutputs(NodeItem &node_item);
   Status LoadKnownShapedSubgraph(ComputeGraph &graph, NodeItem *parent_node_item);

ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc

@@ -18,6 +18,7 @@
 #include "cce/aicpu_engine_struct.h"
 #include "framework/common/debug/ge_log.h"
 #include "framework/common/fmk_error_codes.h"
+#include "common/dump/dump_manager.h"
 #include "common/ge/ge_util.h"
 #include "graph/attr_value.h"
 #include "graph/debug/ge_attr_define.h"
@@ -110,15 +111,6 @@ Status KnownNodeTask::Init(TaskContext &context) {
     GELOGI("KnownNodeTask::Init mem base is %p, size %lu.",
            davinci_model_->GetRuntimeParam().mem_base, davinci_model_->GetRuntimeParam().mem_size);
   }
-  if (!load_flag_) {
-    auto dump_properties = context.GetDumpProperties();
-    if (dump_properties.IsDumpOpen() || dump_properties.IsOpDebugOpen()) {
-      davinci_model_->SetDumpProperties(dump_properties);
-      void *global_step = context.GetExecutionContext()->global_step;
-      davinci_model_->SetKnownShapeGlobalStep(global_step);
-    }
-    load_flag_ = true;
-  }
   GE_CHK_STATUS_RET(ModelManager::GetInstance()->DestroyAicpuKernel(davinci_model_->GetSessionId(),
                                                                     davinci_model_->Id(), davinci_model_->SubModelId()),
                     "KnownNodeTask::Init destroy aicpu kernel failed.");
@@ -126,13 +118,21 @@ Status KnownNodeTask::Init(TaskContext &context) {
   return SUCCESS;
 }
 
-Status KnownNodeTask::InitDavinciModel() {
+Status KnownNodeTask::InitDavinciModel(const HybridModel &model) {
   GELOGD("[Init][Model] start");
   davinci_model_->InitRuntimeParams();
   GE_CHK_STATUS_RET(davinci_model_->InitVariableMem(), "init variable mem failed");
   int32_t device_id = 0;
   GE_CHK_RT_RET(rtGetDevice(&device_id));
   davinci_model_->SetDeviceId(static_cast<uint32_t>(device_id));
+
+  auto dump_properties = DumpManager::GetInstance().GetDumpProperties(model.GetSessionId());
+  if (dump_properties.IsDumpOpen() || dump_properties.IsOpDebugOpen()) {
+    davinci_model_->SetDumpProperties(dump_properties);
+    void *global_step = model.GetGlobalStep();
+    davinci_model_->SetKnownShapeGlobalStep(global_step);
+  }
+
   GE_CHK_STATUS_RET(DoInitDavinciModel(), "[Init][Model] Failed to init davinci model.");
   GELOGD("[Init][Model] success");
   return SUCCESS;
@@ -180,7 +180,7 @@ Status KnownNodeExecutor::LoadTask(const HybridModel &model, const NodePtr &node
 
   auto known_node_task = MakeShared<KnownNodeTask>(davinci_model);
   GE_CHECK_NOTNULL(known_node_task);
-  GE_CHK_STATUS_RET_NOLOG(known_node_task->InitDavinciModel());
+  GE_CHK_STATUS_RET_NOLOG(known_node_task->InitDavinciModel(model));
   GELOGI("[%s] KnownNodeExecutor::LoadTask success.", node->GetName().c_str());
   task = std::move(known_node_task);
   return SUCCESS;

ge/hybrid/node_executor/compiledsubgraph/known_node_executor.h

@@ -36,7 +36,7 @@ class KnownNodeTask : public NodeTask {
   Status UpdateArgs(TaskContext &context) override;
   Status ExecuteAsync(TaskContext &context, std::function<void()> done_callback) override;
   Status Init(TaskContext &context) override;
-  Status InitDavinciModel();
+  Status InitDavinciModel(const HybridModel &model);
 
  protected:
   virtual Status DoInitDavinciModel();