graphengine/ge/graph/passes/dynamic_single_op_reset_shape_pass.cc

/**
 * Copyright 2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "graph/passes/dynamic_single_op_reset_shape_pass.h"
#include "graph/utils/node_utils.h"
#include "graph/utils/graph_utils.h"
#include "graph/utils/tensor_utils.h"
#include "graph/utils/op_desc_utils.h"
#include "graph/utils/type_utils.h"
#include "graph/debug/ge_attr_define.h"

namespace ge {
Status DynamicSingleOpResetShapePass::Run(ComputeGraphPtr graph) {
  GE_CHECK_NOTNULL(graph);
  for (const auto &node : graph->GetDirectNode()) {
    GE_CHECK_NOTNULL(node->GetOpDesc());
    if (node->GetType() == DATA || node->GetType() == NETOUTPUT) {
      continue;
    }

    bool single_aicpu_unknown = false;
    if (!AttrUtils::GetBool(node->GetOpDesc(), ATTR_DYNAMIC_SHAPE_SINGLE_AICPU, single_aicpu_unknown) ||
        !single_aicpu_unknown) {
      continue;
    }

    auto op_desc = node->GetOpDesc();
    std::vector<int64_t> dynamic_shape_dims = {-2};
    ge::Shape dynamic_shape(dynamic_shape_dims);
    for (size_t i = 0; i < op_desc->GetAllInputsDesc().size(); i++) {
      const auto &input_desc = op_desc->MutableInputDesc(static_cast<uint32_t>(i));
      GE_CHECK_NOTNULL(input_desc);
      input_desc.SetShape(dynamic_shape);
    }
    GELOGD("Reset dynamic aicpu node [%s] shape success!", node->GetName().c_str());
  }

  GELOGD("Reset dynamic aicpu nodes shape of graph [%s] success!", graph->GetName().c_str());
  return SUCCESS;
}
}  // namespace ge