!1204 Bugfix: Tiger online inference change on executor

From: @hugo1
Reviewed-by: 
Signed-off-by:

ge/graph/build/graph_builder.cc

@@ -400,6 +400,10 @@ static Status InsertMemcpyNode(const ComputeGraphPtr &graph, const OutDataAnchor
 }
 
 static Status GenerateTaskForConstant(const std::shared_ptr<ComputeGraph> &graph) {
+  if (graph->GetGraphUnknownFlag()) {
+    GELOGI("Graph %s is unknown graph, ignore gen_task for constant.", graph->GetName().c_str());
+    return SUCCESS;
+  }
   for (auto &node : graph->GetDirectNode()) {
     // CONSTANT not generate task, so insert IDENTITY between CONSTANT and NETOUTPUT
     auto op_desc = node->GetOpDesc();

ge/graph/load/model_manager/model_manager.cc

@@ -298,10 +298,8 @@ Status ModelManager::LoadModelOnline(uint32_t &model_id, const shared_ptr<ge::Ge
     GenModelId(&model_id);
   }
 
-  bool is_shape_unknown = false;
   string model_name = "";
-  GE_CHK_STATUS_RET(ge_root_model->CheckIsUnknownShape(is_shape_unknown), "CheckIsUnknownShape failed, model id:%u",
-                    model_id);
+  bool is_shape_unknown = ge_root_model->GetRootGraph()->GetGraphUnknownFlag();
   if (is_shape_unknown || GetContext().GetHostExecFlag()) {
     return DoLoadHybridModelOnline(model_id, model_name, ge_root_model, listener);
   }

ge/host_kernels/concat_offset_kernel.cc

@@ -33,7 +33,7 @@ const int kNumOne = 1;
 }  // namespace
 Status ConcatOffsetKernel::Compute(const OpDescPtr op_desc_ptr, const vector<ConstGeTensorPtr> &input,
                                    vector<GeTensorPtr> &v_output) {
-  GELOGI("ConcatOffsetKernel in.");
+  GELOGD("ConcatOffsetKernel in");
   if (op_desc_ptr == nullptr) {
     GELOGE(PARAM_INVALID, "input opdesc is nullptr.");
     return PARAM_INVALID;
@@ -41,7 +41,7 @@ Status ConcatOffsetKernel::Compute(const OpDescPtr op_desc_ptr, const vector<Con
   // validate attrs
   int N = 0;
   if (!(AttrUtils::GetInt(op_desc_ptr, "N", N))) {
-    GELOGW("Attr %s does not exist.", "N");
+    GELOGW("Attr %s does not exist", "N");
     return NOT_CHANGED;
   }
   // follow IR def, the first input is concat_dim
@@ -50,7 +50,7 @@ Status ConcatOffsetKernel::Compute(const OpDescPtr op_desc_ptr, const vector<Con
   int32_t concat_dim = *(const_cast<int32_t *>(reinterpret_cast<const int32_t *>(input_0->GetData().data())));
   // validate inputs
   if ((static_cast<int>(input.size()) != (N + kNumOne)) || (input.size() <= kConcatOffsetInputIndexOne)) {
-    GELOGW("The number of input for concat offset must be equal to %d, and must be more than one.", (N + kNumOne));
+    GELOGW("The number of input for concat offset must be equal to %d, and must be more than one", (N + kNumOne));
     return NOT_CHANGED;
   }
 
@@ -61,7 +61,7 @@ Status ConcatOffsetKernel::Compute(const OpDescPtr op_desc_ptr, const vector<Con
     GELOGW("Concat dim is bigger than the size of output_shape.");
     return NOT_CHANGED;
   }
-  GELOGI("Output shape size is %ld", output_size);
+  GELOGI("Output shape size is %ld.", output_size);
   int32_t offset = 0;
   if (output_size < 0) {
     GELOGE(FAILED, "Index is negative.");
@@ -86,7 +86,7 @@ Status ConcatOffsetKernel::Compute(const OpDescPtr op_desc_ptr, const vector<Con
     output_ptr->MutableTensorDesc().SetShape(output_shape);
     GE_IF_BOOL_EXEC(output_ptr->SetData(reinterpret_cast<uint8_t *>(buf.get()),
                                         static_cast<size_t>(sizeof(DT_INT32) * output_size)) != GRAPH_SUCCESS,
-                    GELOGW("set data failed");
+                    GELOGW("set data failed.");
                     return NOT_CHANGED);
     v_output.push_back(output_ptr);
     // caculate offset
@@ -99,7 +99,7 @@ Status ConcatOffsetKernel::Compute(const OpDescPtr op_desc_ptr, const vector<Con
     }
     offset += input_dim;
   }
-  GELOGI("ConcatOffsetKernel success.");
+  GELOGD("ConcatOffsetKernel success");
   return SUCCESS;
 }
 REGISTER_KERNEL(CONCATOFFSET, ConcatOffsetKernel);

ge/host_kernels/gather_v2_kernel.cc

@@ -278,7 +278,7 @@ Status GatherV2Kernel::SaveIndicesByDataType(ConstGeTensorPtr indices_tensor_ptr
     auto indices_ptr = const_cast<int32_t *>(reinterpret_cast<const int32_t *>(indices_tensor_ptr->GetData().data()));
     for (int64_t i = 0; i < indices_shape.GetShapeSize(); i++) {
       if (*(indices_ptr + i) < 0 || *(indices_ptr + i) >= x_shape.GetDim(axis)) {
-        GELOGW("indices %ld value is not in range [0, %ld)", i, x_shape.GetDim(axis));
+        GELOGW("indices %ld value is not in range [0, %ld).", i, x_shape.GetDim(axis));
         return NOT_CHANGED;
       }
       indicates_.push_back(*(indices_ptr + i));
@@ -288,7 +288,7 @@ Status GatherV2Kernel::SaveIndicesByDataType(ConstGeTensorPtr indices_tensor_ptr
     auto indices_ptr = const_cast<int64_t *>(reinterpret_cast<const int64_t *>(indices_tensor_ptr->GetData().data()));
     for (int64_t i = 0; i < indices_shape.GetShapeSize(); i++) {
       if (*(indices_ptr + i) < 0 || *(indices_ptr + i) >= x_shape.GetDim(axis)) {
-        GELOGW("indices %ld value is not in range [0, %ld)", i, x_shape.GetDim(axis));
+        GELOGW("indices %ld value is not in range [0, %ld).", i, x_shape.GetDim(axis));
         return NOT_CHANGED;
       }
       indicates_.push_back(*(indices_ptr + i));
@@ -344,42 +344,42 @@ Status GatherV2Kernel::Check(const OpDescPtr &op_desc_ptr, const vector<ConstGeT
   auto indices_data_type = tensor1->GetTensorDesc().GetDataType();
   bool is_valid_indices_data_type = indices_data_type == DT_INT32 || indices_data_type == DT_INT64;
   if (!is_valid_indices_data_type) {
-    GELOGW("indices datatype must be DT_INT32 or DT_INT64");
+    GELOGW("indices datatype must be DT_INT32 or DT_INT64.");
     return NOT_CHANGED;
   }
   if (indices_shape.GetDimNum() > kMaxIndicatesDims) {
-    GELOGW("indices input only support 0 or 1 dims");
+    GELOGW("indices input only support 0 or 1 dims.");
     return NOT_CHANGED;
   }
   return SUCCESS;
 }
 void GatherV2Kernel::DebugPrint(int64_t axis, const GeShape &x_shape, const GeShape &indices_shape,
                                 const std::vector<int64_t> &y_shape) {
-  GELOGD("GatherV2Kernel axis:%ld x_shape:%zu indices_shape:%zu y_shape:%zu", axis, x_shape.GetDimNum(),
+  GELOGD("GatherV2Kernel axis:%ld x_shape:%zu indices_shape:%zu y_shape:%zu.", axis, x_shape.GetDimNum(),
          indices_shape.GetDimNum(), y_shape.size());
   for (size_t i = 0; i < x_shape.GetDimNum(); i++) {
-    GELOGD("GatherV2Kernel x_shape[%zu]: %ld", i, x_shape.GetDim(i));
+    GELOGD("GatherV2Kernel x_shape[%zu]: %ld.", i, x_shape.GetDim(i));
   }
   for (size_t i = 0; i < indices_shape.GetDimNum(); i++) {
-    GELOGD("GatherV2Kernel indices_shape[%zu]: %ld", i, indices_shape.GetDim(i));
+    GELOGD("GatherV2Kernel indices_shape[%zu]: %ld.", i, indices_shape.GetDim(i));
   }
   for (size_t i = 0; i < y_shape.size(); i++) {
-    GELOGD("GatherV2Kernel y_shape[%zu]: %ld", i, y_shape[i]);
+    GELOGD("GatherV2Kernel y_shape[%zu]: %ld.", i, y_shape[i]);
   }
   for (auto ele : indicates_) {
-    GELOGD("GatherV2Kernel indices:%ld", ele);
+    GELOGD("GatherV2Kernel indices:%ld.", ele);
   }
 }
 
 Status GatherV2Kernel::Compute(const OpDescPtr op_desc_ptr, const vector<ConstGeTensorPtr> &input,
                                vector<GeTensorPtr> &v_output) {
-  GELOGI("Enter GatherV2Kernel Process.");
+  GELOGI("Enter GatherV2Kernel Process");
   Status ret = Check(op_desc_ptr, input, v_output);
   if (ret != SUCCESS) {
-    GELOGW("param check failed.");
+    GELOGW("param check failed");
     return NOT_CHANGED;
   }
-  GELOGI("GatherV2Kernel[%s] start Process.", op_desc_ptr->GetName().c_str());
+  GELOGI("GatherV2Kernel[%s] start Process", op_desc_ptr->GetName().c_str());
   ConstGeTensorPtr tensor0 = input.at(kGatherV2InputIndexZero);
   ConstGeTensorPtr tensor1 = input.at(kGatherV2InputIndexOne);
   ConstGeTensorPtr tensor2 = input.at(kGatherV2InputIndexTwo);
@@ -394,7 +394,7 @@ Status GatherV2Kernel::Compute(const OpDescPtr op_desc_ptr, const vector<ConstGe
   axis = axis >= 0 ? axis : axis + x_shape.GetDimNum();
   // check axis value
   if (axis < 0 || (axis + 1) > static_cast<int64_t>(x_shape.GetDimNum())) {
-    GELOGW("axis is invalid");
+    GELOGW("axis is invalid!");
     return NOT_CHANGED;
   }
   auto indices_data_type = tensor1->GetTensorDesc().GetDataType();
@@ -407,7 +407,7 @@ Status GatherV2Kernel::Compute(const OpDescPtr op_desc_ptr, const vector<ConstGe
   // check input data type
   auto x_data_type = tensor0->GetTensorDesc().GetDataType();
   if (supported_type.find(x_data_type) == supported_type.end()) {
-    GELOGI("GatherV2Kernel does not support this Data type:%s", TypeUtils::DataTypeToSerialString(x_data_type).c_str());
+    GELOGI("GatherV2Kernel does not support this Data type:%s.", TypeUtils::DataTypeToSerialString(x_data_type).c_str());
     return NOT_CHANGED;
   }
   // calc output shape

ge/host_kernels/identity_kernel.cc

@@ -61,4 +61,5 @@ Status IdentityKernel::Compute(const ge::OpDescPtr op_desc, const std::vector<ge
   return SUCCESS;
 }
 REGISTER_KERNEL(IDENTITY, IdentityKernel);
+REGISTER_KERNEL(PLACEHOLDERWITHDEFAULT, IdentityKernel);
 }  // namespace ge

ge/host_kernels/strided_slice_kernel.cc

@@ -84,14 +84,14 @@ void GetOriginStrideVec(const std::vector<ge::ConstGeTensorPtr> &input, vector<i
 }  // namespace
 Status StridedSliceKernel::Compute(const ge::OpDescPtr attr, const std::vector<ge::ConstGeTensorPtr> &input,
                                    vector<ge::GeTensorPtr> &v_output) {
-  GELOGD("StridedSliceKernel in.");
+  GELOGD("StridedSliceKernel in");
   // 1.Check input and attrs
   if (CheckAndGetAttr(attr) != SUCCESS) {
-    GELOGW("Check and get attrs failed.Ignore kernel.");
+    GELOGW("Check and get attrs failed.Ignore kernel");
     return NOT_CHANGED;
   }
   if (CheckInputParam(input) != SUCCESS) {
-    GELOGW("Check input params failed.Ignore kernel.");
+    GELOGW("Check input params failed.Ignore kernel");
     return NOT_CHANGED;
   }
   // 2.Init param with mask attrs.
@@ -120,7 +120,7 @@ Status StridedSliceKernel::Compute(const ge::OpDescPtr attr, const std::vector<g
   auto ret = OpUtils::SetOutputSliceData(data, static_cast<int64_t>(data_size), data_type, input_dims, begin_vec,
                                          output_dims, output_ptr.get(), stride_vec);
   if (ret != SUCCESS) {
-    GELOGE(INTERNAL_ERROR, "SetOutputSliceData failed.");
+    GELOGE(INTERNAL_ERROR, "SetOutputSliceData failed");
     return NOT_CHANGED;
   }
 
@@ -133,7 +133,7 @@ Status StridedSliceKernel::Compute(const ge::OpDescPtr attr, const std::vector<g
   GetOutputDims(final_dim_size, output_dims, v_dims);
   t_d.SetShape(GeShape(v_dims));
   v_output.push_back(output_ptr);
-  GELOGI("StridedSliceKernel success.");
+  GELOGI("StridedSliceKernel success");
   return SUCCESS;
 }
 Status StridedSliceKernel::CheckAndGetAttr(const OpDescPtr &attr) {
@@ -144,7 +144,7 @@ Status StridedSliceKernel::CheckAndGetAttr(const OpDescPtr &attr) {
   // Get all op attr value of strided_slice
   for (auto &attr_2_value : attr_value_map_) {
     if (!AttrUtils::GetInt(attr, attr_2_value.first, attr_2_value.second)) {
-      GELOGE(PARAM_INVALID, "Get %s attr failed.", attr_2_value.first.c_str());
+      GELOGE(PARAM_INVALID, "Get %s attr failed", attr_2_value.first.c_str());
       return PARAM_INVALID;
     }
   }
@@ -182,7 +182,7 @@ Status StridedSliceKernel::CheckInputParam(const std::vector<ConstGeTensorPtr> &
     return PARAM_INVALID;
   }
   if (kIndexNumberType.find(begin_tensor_desc.GetDataType()) == kIndexNumberType.end()) {
-    GELOGW("Data type of StridedSlice OP(begin,end,strides) must be int32 or int64.");
+    GELOGW("Data type of StridedSlice OP(begin,end,strides) must be int32 or int64");
     return PARAM_INVALID;
   }
 
@@ -250,7 +250,7 @@ Status StridedSliceKernel::InitParamWithAttrs(const std::vector<ConstGeTensorPtr
       end_i = x_dims.at(i);
       stride_i = 1;
     }
-    GELOGD("Before mask calculate. Begin is : %ld\t,end is : %ld\t stride is : %ld\t x_dim_i is : %ld.",
+    GELOGD("Before mask calculate. Begin is : %ld\t,end is : %ld\t stride is : %ld\t x_dim_i is : %ld",
            begin_i, end_i, stride_i, x_dims.at(i));
     auto ret = MaskCal(i, begin_i, end_i, x_dims.at(i));
     if (ret != SUCCESS) {
@@ -258,7 +258,7 @@ Status StridedSliceKernel::InitParamWithAttrs(const std::vector<ConstGeTensorPtr
       return NOT_CHANGED;
     }
     int64_t dim_final;
-    GELOGD("Before stride calculate. Begin is : %ld\t,end is : %ld\t stride is : %ld\t x_dim_i is : %ld.",
+    GELOGD("Before stride calculate. Begin is : %ld\t,end is : %ld\t stride is : %ld\t x_dim_i is : %ld",
            begin_i, end_i, stride_i, x_dims.at(i));
     (void) StrideCal(x_dims.at(i), begin_i, end_i, stride_i, dim_final);
     output_dims.push_back(dim_final);

ge/hybrid/model/hybrid_model.h

@@ -154,6 +154,7 @@ class HybridModel {
   uint32_t model_id_ = 0;
   uint8_t *var_mem_base_ = nullptr;
   std::unique_ptr<TensorBuffer> weight_buffer_;
+  std::map<string, std::unique_ptr<TensorBuffer>> weight_buffer_map_;
   RuntimeParam root_runtime_param_;
   string om_name_;
 };

ge/hybrid/model/hybrid_model_builder.cc

@@ -1012,70 +1012,65 @@ Status HybridModelBuilder::InitVariableTensors() {
 
 Status HybridModelBuilder::InitWeights() {
   // For constant in root graph
-  const auto &root_graph = ge_root_model_->GetRootGraph();
-  const auto &subgraph_models = ge_root_model_->GetSubgraphInstanceNameToModel();
-  auto iter = subgraph_models.find(root_graph->GetName());
-  if (iter == subgraph_models.end()) {
-    GELOGD("Root graph model not found");
-    return SUCCESS;
-  }
+  for (const auto &subgraph_model : ge_root_model_->GetSubgraphInstanceNameToModel()) {
+    const auto &weight_buffer = subgraph_model.second->GetWeight();
+    if (weight_buffer.GetSize() == 0) {
+      GELOGD("weight is empty");
+      return SUCCESS;
+    }
 
-  auto &root_model = iter->second;
-  const auto &weight_buffer = root_model->GetWeight();
-  if (weight_buffer.GetSize() == 0) {
-    GELOGD("weight is empty");
-    return SUCCESS;
-  }
+    auto allocator = NpuMemoryAllocator::GetAllocator();
+    GE_CHECK_NOTNULL(allocator);
+    auto sub_weight_buffer = TensorBuffer::Create(allocator, weight_buffer.size());
+    GE_CHECK_NOTNULL(sub_weight_buffer);
+    auto weight_base = reinterpret_cast<uint8_t *>(sub_weight_buffer->GetData());
+    GE_CHK_RT_RET(rtMemcpy(weight_base,
+                           sub_weight_buffer->GetSize(),
+                           weight_buffer.GetData(),
+                           weight_buffer.GetSize(),
+                           RT_MEMCPY_HOST_TO_DEVICE));
+
+    GELOGI("Init weight mem successfully, weight base %p, weight size = %zu",
+           weight_base,
+           sub_weight_buffer->GetSize());
+    auto root_graph = GraphUtils::GetComputeGraph(subgraph_model.second->GetGraph());
+    hybrid_model_.weight_buffer_map_.emplace(root_graph->GetName(),std::move(sub_weight_buffer));
+    for (auto &node : root_graph->GetDirectNode()) {
+      if (node->GetType() != CONSTANT) {
+        continue;
+      }
 
-  auto allocator = NpuMemoryAllocator::GetAllocator();
-  GE_CHECK_NOTNULL(allocator);
-  hybrid_model_.weight_buffer_ = TensorBuffer::Create(allocator, weight_buffer.size());
-  GE_CHECK_NOTNULL(hybrid_model_.weight_buffer_);
-  auto weight_base = reinterpret_cast<uint8_t *>(hybrid_model_.weight_buffer_->GetData());
-  GE_CHK_RT_RET(rtMemcpy(weight_base,
-                         hybrid_model_.weight_buffer_->GetSize(),
-                         weight_buffer.GetData(),
-                         weight_buffer.GetSize(),
-                         RT_MEMCPY_HOST_TO_DEVICE));
-
-  GELOGI("Init weight mem successfully, weight base %p, weight size = %zu",
-         weight_base,
-         hybrid_model_.weight_buffer_->GetSize());
-  for (auto &node : root_graph->GetDirectNode()) {
-    if (node->GetType() != CONSTANT) {
-      continue;
-    }
+      auto op_desc = node->GetOpDesc();
+      auto v_weights = ModelUtils::GetWeights(op_desc);
+      if (v_weights.empty()) {
+        GELOGE(INTERNAL_ERROR, "[%s] Constant has no value", node->GetName().c_str());
+        return INTERNAL_ERROR;
+      }
+      auto *ge_tensor = const_cast<GeTensor *>(v_weights[0].get());
+      GE_CHECK_NOTNULL(ge_tensor);
+      const GeTensorDesc &tensor_desc = ge_tensor->GetTensorDesc();
+      int64_t tensor_size = 0;
+      GE_CHK_GRAPH_STATUS_RET(TensorUtils::GetSize(*op_desc->MutableOutputDesc(0), tensor_size),
+                              "[%s] Failed to get tensor size",
+                              node->GetName().c_str());
+      int64_t data_offset = 0;
+      GE_CHK_GRAPH_STATUS_RET(TensorUtils::GetDataOffset(tensor_desc, data_offset),
+                              "[%s] Failed to get data offset",
+                              node->GetName().c_str());
+      GELOGD("[%s] Start to init Constant node [%s], size = %ld, offset = %ld",
+             GetGraphName(),
+             node->GetName().c_str(),
+             tensor_size,
+             data_offset);
 
-    auto op_desc = node->GetOpDesc();
-    auto v_weights = ModelUtils::GetWeights(op_desc);
-    if (v_weights.empty()) {
-      GELOGE(INTERNAL_ERROR, "[%s] Constant has no value", node->GetName().c_str());
-      return INTERNAL_ERROR;
+      auto tensor_buffer = TensorBuffer::Create(weight_base + data_offset, tensor_size);
+      GE_CHECK_NOTNULL(tensor_buffer);
+      std::unique_ptr<TensorValue> constant_tensor(new (std::nothrow)TensorValue(std::move(tensor_buffer)));
+      GE_CHECK_NOTNULL(constant_tensor);
+      constant_tensor->SetName("Constant_" + op_desc->GetName());
+      hybrid_model_.constant_tensors_.emplace(node, std::move(constant_tensor));
+      GELOGD("[%s] Constant node [%s] added, size = %ld", GetGraphName(), node->GetName().c_str(), tensor_size);
     }
-    auto *ge_tensor = const_cast<GeTensor *>(v_weights[0].get());
-    GE_CHECK_NOTNULL(ge_tensor);
-    const GeTensorDesc &tensor_desc = ge_tensor->GetTensorDesc();
-    int64_t tensor_size = 0;
-    GE_CHK_GRAPH_STATUS_RET(TensorUtils::GetSize(*op_desc->MutableOutputDesc(0), tensor_size),
-                            "[%s] Failed to get tensor size",
-                            node->GetName().c_str());
-    int64_t data_offset = 0;
-    GE_CHK_GRAPH_STATUS_RET(TensorUtils::GetDataOffset(tensor_desc, data_offset),
-                            "[%s] Failed to get data offset",
-                            node->GetName().c_str());
-    GELOGD("[%s] Start to init Constant node [%s], size = %ld, offset = %ld",
-           GetGraphName(),
-           node->GetName().c_str(),
-           tensor_size,
-           data_offset);
-
-    auto tensor_buffer = TensorBuffer::Create(weight_base + data_offset, tensor_size);
-    GE_CHECK_NOTNULL(tensor_buffer);
-    std::unique_ptr<TensorValue> constant_tensor(new (std::nothrow)TensorValue(std::move(tensor_buffer)));
-    GE_CHECK_NOTNULL(constant_tensor);
-    constant_tensor->SetName("Constant_" + op_desc->GetName());
-    hybrid_model_.constant_tensors_.emplace(node, std::move(constant_tensor));
-    GELOGD("[%s] Constant node [%s] added, size = %ld", GetGraphName(), node->GetName().c_str(), tensor_size);
   }
   return SUCCESS;
 }

ge/hybrid/node_executor/task_context.cc

@@ -236,7 +236,7 @@ Status TaskContext::AllocateOutput(int index,
            ref_node->GetName().c_str(),
            ref_node->GetType().c_str());
 
-    TensorValue *ref_tensor = execution_context_->model->GetVariable(ref_node->GetName());
+    TensorValue *ref_tensor = execution_context_->model->GetTensor(ref_node);
     GE_CHECK_NOTNULL(ref_tensor);
     outputs_start_[index] = *ref_tensor;
   } else {

tests/ut/ge/hybrid/ge_hybrid_unittest.cc

@@ -15,8 +15,8 @@
  */
 
 #include <gtest/gtest.h>
+#include <gmock/gmock.h>
 #include <vector>
-
 #include "runtime/rt.h"
 
 #define protected public
@@ -25,7 +25,6 @@
 #include "hybrid/model/hybrid_model.h"
 #include "model/ge_model.h"
 #include "model/ge_root_model.h"
-
 #include "hybrid/node_executor/aicore/aicore_op_task.h"
 #include "framework/common/taskdown_common.h"
 #include "framework/common/debug/log.h"
@@ -33,7 +32,10 @@
 #include "hybrid/executor/hybrid_execution_context.h"
 #include "hybrid/node_executor/aicore/aicore_task_builder.h"
 #include "graph/load/model_manager/tbe_handle_store.h"
+#include "graph/manager/graph_mem_allocator.h"
+#include "hybrid/common/npu_memory_allocator.h"
 #include "graph/types.h"
+#include "graph/utils/tensor_utils.h"
 
 #undef private
 #undef protected
@@ -43,6 +45,7 @@ using namespace testing;
 using namespace ge;
 using namespace hybrid;
 
+
 class UtestGeHybrid : public testing::Test {
  protected:
   void SetUp() {}
@@ -204,4 +207,39 @@ TEST_F(UtestGeHybrid, index_taskdefs_success) {
   HybridModelBuilder hybrid_model_builder(hybrid_model);
 
   ASSERT_EQ(hybrid_model_builder.IndexTaskDefs(graph, ge_model), SUCCESS);
+}
+
+TEST_F(UtestGeHybrid, init_weight_success) {
+  NpuMemoryAllocator::allocators_.emplace(make_pair(0, nullptr));
+  // make graph with sub_graph
+  ComputeGraphPtr graph = std::make_shared<ComputeGraph>("root_graph");
+  OpDescPtr op_desc = CreateOpDesc("if", IF);
+  NodePtr node = graph->AddNode(op_desc);
+  // make sub graph
+  ComputeGraphPtr sub_graph = std::make_shared<ComputeGraph>("if_sub_graph");
+  OpDescPtr const_op_desc = CreateOpDesc("const", CONSTANT);
+  vector<int64_t> dims_vec_0 = {2, 1, 4, 1, 2};
+  vector<int32_t> data_vec_0 = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16};
+  GeTensorDesc tensor_desc_0(GeShape(dims_vec_0), FORMAT_NCHW, DT_INT32);
+  (void)TensorUtils::SetRealDimCnt(tensor_desc_0, dims_vec_0.size());
+  ConstGeTensorPtr constTensor_0 =
+    std::make_shared<GeTensor>(tensor_desc_0, (uint8_t *)&data_vec_0[0], data_vec_0.size() * sizeof(int32_t));
+  AttrUtils::SetTensor(const_op_desc, ge::ATTR_NAME_WEIGHTS, constTensor_0);
+  const_op_desc->AddOutputDesc(tensor_desc_0);
+  NodePtr const_node = sub_graph->AddNode(const_op_desc);
+  graph->AddSubgraph("sub", sub_graph);
+
+  GeRootModelPtr ge_root_model = make_shared<GeRootModel>(graph);
+  GeModelPtr ge_sub_model = make_shared<GeModel>();
+  //Buffer weight_buffer = Buffer(128,0);
+  //ge_sub_model->SetWeight(weight_buffer);
+  ge_root_model->SetSubgraphInstanceNameToModel("sub",ge_sub_model);
+  HybridModel hybrid_model(ge_root_model);
+  HybridModelBuilder hybrid_model_builder(hybrid_model);
+  auto ret = hybrid_model_builder.InitWeights();
+  ASSERT_EQ(ret,SUCCESS);
+  Buffer weight_buffer = Buffer(128,0);
+  ge_sub_model->SetWeight(weight_buffer);
+  ret = hybrid_model_builder.InitWeights();
+  ASSERT_EQ(ret,PARAM_INVALID);
 }