add inner offset for var addr judgement and update

ge/graph/build/memory/graph_mem_assigner.cc

@@ -1545,6 +1545,13 @@ ge::Status GraphMemoryAssigner::UpdateOpInputOffset(const NodePtr &node, vector<
           input_offset = origin_input_list[valid_input_index];
         } else {
           // hbm input_offset = original input_offset + output_offset
+          if (origin_input_list[valid_input_index] != 0 && tmp_op_desc->GetSubgraphInstanceNames().size() > 0) {
+            std::string error = "Node" + FmtToStr(tmp_op_desc->GetName()) +
+                                +" has subgraphs which is conflict with has origin_input_list" +
+                                FmtToStr(origin_input_list[valid_input_index]);
+            GE_ERRORLOG_AND_ERRORMSG(ge::FAILED, error.c_str());
+            return ge::FAILED;
+          }
           input_offset = origin_input_list[valid_input_index] + output_list.at(out_index);
           (void)ge::AttrUtils::SetInt(tmp_op_desc->MutableInputDesc(anchor->GetIdx()), ATTR_NAME_INNER_OFFSET,
                                       origin_input_list[valid_input_index] + inner_offset);

ge/graph/build/task_generator.cc

@@ -187,33 +187,61 @@ Status TaskGenerator::AddModelTaskToModel(const ModelTaskDef &model_task_def, ui
 }
 
 Status TaskGenerator::UpdateOpIsVarAttr(const OpDescPtr &op_desc, uint64_t session_id) {
-  vector<int64_t> input_offsets = op_desc->GetInputOffset();
   GELOGD("Update is var attr, node[name:%s(%s), id:%ld, stream_id:%ld].", op_desc->GetName().c_str(),
          op_desc->GetType().c_str(), op_desc->GetId(), op_desc->GetStreamId());
+  // input
+  vector<int64_t> input_offsets = op_desc->GetInputOffset();
   if (!(input_offsets.empty())) {
     vector<bool> input_var;
-    for (int64_t input : input_offsets) {
-      input_var.push_back(VarManager::Instance(session_id)->IsVarAddr(input));
+    int64_t valid_input_index = 0;
+    for (uint32_t i = 0; i < op_desc->GetAllInputsSize(); i++) {
+      vector<int64_t> output_list;
+      auto input_tensor_desc = op_desc->MutableInputDesc(i);
+      if (input_tensor_desc == nullptr) {
+        continue;
+      }
+      if (valid_input_index >= input_offsets.size()) {
+        break;
+      }
+      int64_t inner_offset = 0;
+      (void)ge::AttrUtils::GetInt(input_tensor_desc, ATTR_NAME_INNER_OFFSET, inner_offset);
+      GELOGD("Node[%s] input[%u] has inner_offset[%ld]", op_desc->GetName().c_str(), i, inner_offset);
+      input_var.push_back(VarManager::Instance(session_id)->IsVarAddr(input_offsets[valid_input_index] - inner_offset));
+      valid_input_index++;
     }
     GE_CHK_BOOL_EXEC(AttrUtils::SetListBool(op_desc, kIsInputVar, input_var),
                      REPORT_INNER_ERROR("E19999", "Set Attr:%s fail for op:%s(%s)", kIsInputVar,
                                         op_desc->GetName().c_str(), op_desc->GetType().c_str());
-                     GELOGE(FAILED, "[Set][Attr] %s fail for op:%s(%s)", kIsInputVar,
-                            op_desc->GetName().c_str(), op_desc->GetType().c_str());
+                     GELOGE(FAILED, "[Set][Attr] %s fail for op:%s(%s)", kIsInputVar, op_desc->GetName().c_str(),
+                            op_desc->GetType().c_str());
                      return FAILED);
   }
-
+  // output
   vector<int64_t> output_offsets = op_desc->GetOutputOffset();
   if (!(output_offsets.empty())) {
     vector<bool> output_var;
-    for (int64_t output : output_offsets) {
-      output_var.push_back(VarManager::Instance(session_id)->IsVarAddr(output));
+    int64_t valid_output_index = 0;
+    for (uint32_t i = 0; i < op_desc->GetAllOutputsDescSize(); i++) {
+      vector<int64_t> output_list;
+      auto output_tensor_desc = op_desc->MutableOutputDesc(i);
+      if (output_tensor_desc == nullptr) {
+        continue;
+      }
+      if (valid_output_index >= output_offsets.size()) {
+        break;
+      }
+      int64_t inner_offset = 0;
+      (void)ge::AttrUtils::GetInt(output_tensor_desc, ATTR_NAME_INNER_OFFSET, inner_offset);
+      GELOGD("Node[%s] output[%u] has inner_offset[%ld]", op_desc->GetName().c_str(), i, inner_offset);
+      output_var.push_back(
+        VarManager::Instance(session_id)->IsVarAddr(output_offsets[valid_output_index] - inner_offset));
+      valid_output_index++;
     }
     GE_CHK_BOOL_EXEC(AttrUtils::SetListBool(op_desc, kIsOutputVar, output_var),
                      REPORT_INNER_ERROR("E19999", "Set Attr:%s fail for op:%s(%s)", kIsOutputVar,
                                         op_desc->GetName().c_str(), op_desc->GetType().c_str());
-                     GELOGE(FAILED, "[Set][Attr] %s fail for op:%s(%s)", kIsOutputVar,
-                            op_desc->GetName().c_str(), op_desc->GetType().c_str());
+                     GELOGE(FAILED, "[Set][Attr] %s fail for op:%s(%s)", kIsOutputVar, op_desc->GetName().c_str(),
+                            op_desc->GetType().c_str());
                      return FAILED);
   }
   return SUCCESS;

tests/ut/ge/graph/build/task_generator_unittest.cc

@@ -30,13 +30,18 @@
 #define protected public
 #define private public
 #include "graph/build/task_generator.h"
+#include "graph/manager/graph_mem_manager.h"
+#include "graph/manager/graph_var_manager.h"
 #undef protected
 #undef private
 
 using namespace std;
 using namespace testing;
 using namespace ge;
-
+namespace {
+const char *const kIsInputVar = "INPUT_IS_VAR";
+const char *const kIsOutputVar = "OUTPUT_IS_VAR";
+}
 class UtestTaskGeneratorTest : public testing::Test {
  public:
   ge::ComputeGraphPtr BuildGraphFpProfiling() {
@@ -63,6 +68,31 @@ class UtestTaskGeneratorTest : public testing::Test {
     builder.AddControlEdge(addn1, netoutput);
     return builder.GetGraph();
   }
+  ge::ComputeGraphPtr BuildGraphWithVar(int64_t session_id) {
+    // init
+    MemManager::Instance().Initialize(std::vector<rtMemType_t>({RT_MEMORY_HBM}));
+    VarManager::Instance(session_id)->Init(0, 0, 0, 0);
+    ge::ut::GraphBuilder builder("graph");
+    auto var_input = builder.AddNode("var", "Variable", 1, 1);
+    auto const_input = builder.AddNode("const", "Const", 1, 1);
+    auto assign = builder.AddNode("assgin", "Assign", 2, 1);
+    // add link
+    builder.AddDataEdge(var_input, 0, assign, 0);
+    builder.AddDataEdge(const_input, 0, assign, 1);
+    // set offset
+    var_input->GetOpDesc()->SetOutputOffset({10000});
+    const_input->GetOpDesc()->SetOutputOffset({1000});
+    assign->GetOpDesc()->SetInputOffset({10100, 1000});
+    assign->GetOpDesc()->SetOutputOffset({10100});
+    // set inner offset
+    int64_t inner_offset = 100;
+    ge::AttrUtils::SetInt(assign->GetOpDesc()->MutableInputDesc(0), ATTR_NAME_INNER_OFFSET, inner_offset);
+    ge::AttrUtils::SetInt(assign->GetOpDesc()->MutableOutputDesc(0), ATTR_NAME_INNER_OFFSET, inner_offset);
+    // add var addr
+    VarManager::Instance(session_id)->var_resource_->var_offset_map_.emplace(10000, RT_MEMORY_HBM);
+
+    return builder.GetGraph();
+  }
 
  protected:
   void SetUp() {}
@@ -86,3 +116,25 @@ TEST_F(UtestTaskGeneratorTest, FindLastBpFromBpNode) {
   // netoutput has no data input, return default value 0
   EXPECT_EQ(task_generator.FindLastBpFromBpNode(graph, net_output), 0);
 }
+
+TEST_F(UtestTaskGeneratorTest, UpdateOpIsVarAttr) {
+  int64_t session_id = 0;
+  ge::ComputeGraphPtr graph = BuildGraphWithVar(session_id);
+  graph->SetSessionID(session_id);
+  TaskGenerator task_generator(nullptr, 0);
+  auto assign = graph->FindNode("assgin");
+  task_generator.UpdateOpIsVarAttr(assign->GetOpDesc(), session_id);
+  // input
+  vector<bool> input_var;
+  AttrUtils::GetListBool(assign->GetOpDesc(), kIsInputVar, input_var);
+  EXPECT_EQ(input_var.size(), 2);
+  EXPECT_EQ(input_var[0], true);
+  EXPECT_EQ(input_var[1], false);
+  // output
+  vector<bool> output_var;
+  AttrUtils::GetListBool(assign->GetOpDesc(), kIsOutputVar, output_var);
+  EXPECT_EQ(output_var.size(), 1);
+  EXPECT_EQ(output_var[0], true);
+
+  MemManager::Instance().Finalize();
+}