MindSpore-GraphEngine/ge/host_kernels/constant_of_shape_kernel.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 "host_kernels/constant_of_shape_kernel.h"

#include <memory>
#include <set>

#include "common/debug/log.h"
#include "common/fp16_t.h"
#include "common/op/ge_op_utils.h"
#include "framework/common/debug/ge_log.h"
#include "host_kernels/kernel_utils.h"
#include "graph/utils/type_utils.h"
#include "inc/kernel_factory.h"
#include "framework/common/types.h"

namespace ge {
namespace {
const size_t kAddInputSize = 1;
const size_t kAddOutputSize = 1;
}  // namespace

Status ConstantOfShapeKernel::CheckParam(
    const OpDescPtr &op_desc_ptr, const std::vector<ConstGeTensorPtr> &input) {

  GE_CHECK_NOTNULL(op_desc_ptr);
  // check how many inputs
  if ((input.size() != kAddInputSize)
      || (op_desc_ptr->GetOutputsSize() != kAddOutputSize)) {
    GELOGE(PARAM_INVALID,
           "Input number must be [%zu], output number must be [%zu].",
           kAddInputSize, kAddOutputSize);
    return PARAM_INVALID;
  }
  GE_CHECK_NOTNULL(input[0]);
  // input vector elements must be 1-D or empty tensor
  auto input_shape = input[0]->GetTensorDesc().GetShape();
  if (input_shape.GetShapeSize() == 0) {
    GELOGD("Input 0 is empty tensor.");
  } else {
    size_t dim_num = input_shape.GetDimNum();
    if (dim_num != 1) {
      GELOGE(PARAM_INVALID, "Shape input must be a 1-D tensor, but got [%zu]",
              dim_num);
      return PARAM_INVALID;
    }
  }
  return SUCCESS;
}

Status ConstantOfShapeKernel::Compute(const ge::OpDescPtr op_desc_ptr,
                                      const vector<ge::ConstGeTensorPtr> &input,
                                      vector<ge::GeTensorPtr> &v_output) {
  GELOGI("ConstantOfShapeKernel in.");
  if (CheckParam(op_desc_ptr, input) != SUCCESS) {
    return NOT_CHANGED;
  }

  GeShape out_shape;
  int64_t fill_size = 1;
  Status ret = PARAM_INVALID;

  ConstGeTensorPtr value = MakeShared<const GeTensor>();
  GE_CHECK_NOTNULL(value);
  GeTensorPtr output_ptr = MakeShared<GeTensor>(op_desc_ptr->GetOutputDesc(0));
  GE_CHECK_NOTNULL(output_ptr);

  ConstGeTensorPtr dims = input.at(0);
  if (dims->GetData().size() == 0) {
    GELOGI("Input 0 is empty tensor, then output 0 is a scalar.");
    out_shape = GeShape();
  } else {
    std::vector<int64_t> vec_dim;
    GE_RETURN_IF_ERROR(KernelUtils::CalcDims<int64_t>(dims, vec_dim, fill_size));
    out_shape = GeShape(vec_dim);
  }

  DataType data_type;
  if (!AttrUtils::GetTensor(op_desc_ptr, "value", value)) {
    GELOGE(FAILED, "Get Attr value failed.");
    return FAILED;
  }
  data_type = value->GetTensorDesc().GetDataType();

  ret = PARAM_INVALID;
  switch (data_type) {
#define CASE(dtype, type)                                              \
  case dtype:                                                          \
    ret = KernelUtils::GenData(                                        \
            fill_size,                                                 \
            *reinterpret_cast<const type *>(value->GetData().data()),  \
            output_ptr);                                               \
    break;
    CASE(DT_FLOAT, float)
    CASE(DT_FLOAT16, fp16_t)
    CASE(DT_INT8, int8_t)
    CASE(DT_INT16, int16_t)
    CASE(DT_UINT16, uint16_t)
    CASE(DT_UINT8, uint8_t)
    CASE(DT_INT32, int32_t)
    CASE(DT_INT64, int64_t)
    CASE(DT_UINT32, uint32_t)
    CASE(DT_UINT64, uint64_t)
    CASE(DT_BOOL, bool)
    CASE(DT_DOUBLE, double)
#undef CASE
    default:
      GELOGE(PARAM_INVALID, "Invalid data type: [%s]",
             TypeUtils::DataTypeToSerialString(data_type).c_str());
      break;
  }
  if (ret != SUCCESS) {
    GELOGE(ret, "GenData failed, data_type: [%s]",
           TypeUtils::DataTypeToSerialString(data_type).c_str());
    return ret;
  }
  output_ptr->MutableTensorDesc().SetShape(out_shape);
  output_ptr->MutableTensorDesc().SetDataType(data_type);
  v_output.push_back(output_ptr);

  GELOGI("ConstantOfShapeKernel success.");
  return SUCCESS;
}

REGISTER_KERNEL(CONSTANTOFSHAPE, ConstantOfShapeKernel);
}  // namespace ge