mindspore2022/mindspore/ccsrc/transform/op_adapter.h

892 lines
32 KiB
C++

/**
* Copyright 2019 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.
*/
#ifndef TRANSFORM_OP_ADAPTER_H_
#define TRANSFORM_OP_ADAPTER_H_
#include <memory>
#include <vector>
#include <string>
#include <unordered_map>
#include "transform/op_adapter_util.h"
#include "utils/utils.h"
namespace mindspore {
namespace transform {
static uint32_t CustomInferFunc(const Operator &) { return 0; }
template <typename T>
class OpAdapter : public BaseOpAdapter {
public:
using OpType = T;
OpAdapter() {}
explicit OpAdapter(const ExtraAttr &extra_attr) : extra_attr_(extra_attr) {}
~OpAdapter() override {}
bool IsCustomOp(const OperatorPtr &op) {
MS_EXCEPTION_IF_NULL(op);
auto it = cus_input_map_.find(op->GetOpType());
if (it == cus_input_map_.end()) {
return false;
}
return true;
}
Status GenerateCustomOpInputMap(const CusOperatorPtr &op, const PrimitivePtr &prim) {
MS_EXCEPTION_IF_NULL(op);
MS_EXCEPTION_IF_NULL(prim);
// Create the map of custom op from input index to input name.
std::unordered_map<int, std::string> input_map;
auto value = prim->GetAttr("input_names");
if (value == nullptr) {
cus_output_map_[prim->name()] = input_map;
return NOT_FOUND;
}
auto input_names = GetValue<const std::vector<std::string>>(value);
for (size_t i = 0; i < input_names.size(); ++i) {
// input_map begin form 1
input_map[i + 1] = input_names[i];
op->CustomInputRegister(input_names[i]);
}
if (cus_input_map_.find(prim->name()) == cus_input_map_.end()) {
cus_input_map_[prim->name()] = input_map;
}
return SUCCESS;
}
Status GenerateCustomOpOutputMap(const CusOperatorPtr &op, const PrimitivePtr &prim) {
MS_EXCEPTION_IF_NULL(op);
MS_EXCEPTION_IF_NULL(prim);
// Create the map of custom op from output index to output name.
std::unordered_map<int, std::string> output_map;
auto value = prim->GetAttr("output_names");
if (value == nullptr) {
// generate a empty output_map for it
cus_output_map_[prim->name()] = output_map;
return NOT_FOUND;
}
auto output_names = GetValue<const std::vector<std::string>>(value);
for (size_t i = 0; i < output_names.size(); ++i) {
// output_map begin form 0
output_map[i] = output_names[i];
op->CustomOutputRegister(output_names[i]);
}
if (cus_output_map_.find(prim->name()) == cus_output_map_.end()) {
cus_output_map_[prim->name()] = output_map;
}
return SUCCESS;
}
// Convert ME UserCustom AnfNode to GE CustomOp. And set it's attrs.
OperatorPtr GenerateCustomOp(const AnfNodePtr anf) {
MS_EXCEPTION_IF_NULL(anf);
auto node = anf->cast<CNodePtr>();
if (node == nullptr) {
return nullptr;
}
if (node->inputs().empty()) {
MS_LOG(EXCEPTION) << "length of node inputs is empty";
}
auto prim = GetValueNode<PrimitivePtr>(node->inputs()[0]);
MS_EXCEPTION_IF_NULL(prim);
auto op = std::make_shared<ge::CustomOperator>(node->fullname_with_scope(), prim->name());
if (GenerateCustomOpInputMap(op, prim) != SUCCESS) {
MS_LOG(WARNING) << "Custom op node has no input_names, op[" << prim->name() << "].";
}
if (GenerateCustomOpOutputMap(op, prim) != SUCCESS) {
MS_LOG(WARNING) << "Custom op node has no output_names, op[" << prim->name() << "].";
}
op->CustomInferFuncRegister(CustomInferFunc);
return op;
}
OperatorPtr GenerateNormalOp(const AnfNodePtr &anf) {
OperatorPtr op = nullptr;
// There are duplicate names in ANF graph, do not assign ANF node name to GE
// GE will generate unique name automatically
if (anf != nullptr && anf->fullname_with_scope() != "") {
MS_LOG(DEBUG) << anf->fullname_with_scope();
op = std::make_shared<T>(anf->fullname_with_scope());
} else {
MS_LOG(DEBUG) << "no fullname_with_scope";
op = std::make_shared<T>();
}
// set dynamic output num if op use DYNAMIC_OUTPUT
if ((op != nullptr) && (!dyn_output_map_.empty()) && (anf != nullptr)) {
TypePtr type = anf->Type();
if (type == nullptr) {
MS_LOG(EXCEPTION) << "Dynamic output node:" << op->GetName() << "'s Type is a nullptr!";
}
size_t num = type->isa<Tuple>() ? (type->cast<std::shared_ptr<Tuple>>()->size()) : 1;
MS_LOG(INFO) << "create_dyn_output for node:" << anf->ToString() << ", type:" << type->ToString()
<< ", num:" << num;
dyn_output_map_.begin()->second.create_dyn_output(op, static_cast<unsigned int>(num));
}
return op;
}
OperatorPtr generate(const AnfNodePtr &anf) override {
OperatorPtr op = nullptr;
if (IsCustomCNode(anf)) {
op = GenerateCustomOp(anf);
} else {
op = GenerateNormalOp(anf);
}
return op;
}
OperatorPtr generate(const std::string &op_name) override { return std::make_shared<T>(op_name); }
const std::unordered_map<int, InputDesc> &getInputMap() override { return input_map_; }
const std::unordered_map<unsigned int, AttrDesc> &getInputAttrMap() override { return input_attr_map_; }
const std::unordered_map<int, DynInputDesc> &getDynInputMap() override { return dyn_input_map_; }
const std::unordered_map<int, OutputDesc> &getOutputMap() override { return output_map_; }
Status SetCustomOpInput(const CusOperatorPtr &op, int index, const OperatorPtr &input) {
MS_EXCEPTION_IF_NULL(op);
MS_EXCEPTION_IF_NULL(input);
auto it = cus_input_map_.find(op->GetOpType());
if (it == cus_input_map_.end()) {
return NOT_FOUND;
}
std::unordered_map<int, std::string> &input_map = it->second;
if ((input_map.find(index) != input_map.end())) {
MS_LOG(DEBUG) << "Link op " << input->GetName() << " to " << op->GetName() << ":" << input_map[index];
(void)op->SetInput(input_map[index], *input);
return SUCCESS;
}
return NOT_FOUND;
}
Status SetNormalOpInput(const OperatorPtr &op, int index, const OperatorPtr &input) {
MS_EXCEPTION_IF_NULL(op);
auto it = input_map_.find(index);
if (it != input_map_.end()) {
MS_EXCEPTION_IF_NULL(input);
MS_LOG(DEBUG) << "Link op " << input->GetName() << " to " << op->GetName() << ":" << it->second.name;
it->second.set_op(op, input);
return SUCCESS;
}
return NOT_FOUND;
}
int setInput(const OperatorPtr &op, int index, const OperatorPtr &input) override {
if (IsCustomOp(op)) {
auto cus_op = std::dynamic_pointer_cast<CustomOperator>(op);
return static_cast<int>(SetCustomOpInput(cus_op, index, input));
} else {
return static_cast<int>(SetNormalOpInput(op, index, input));
}
}
Status SetCustomOpInput(const CusOperatorPtr &op, int index, const OutHandler &handle) {
MS_EXCEPTION_IF_NULL(op);
auto it = cus_input_map_.find(op->GetOpType());
if (it == cus_input_map_.end()) {
return NOT_FOUND;
}
std::unordered_map<int, std::string> &input_map = it->second;
if ((handle.op != nullptr) && (input_map.find(index) != input_map.end())) {
if (handle.out.empty()) {
MS_LOG(DEBUG) << "Link op " << handle.op->GetName() << " to " << op->GetName() << ":" << input_map[index];
(void)op->SetInput(input_map[index], *(handle.op));
} else {
MS_LOG(DEBUG) << "Link op " << handle.op->GetName() << ":" << handle.out << " to " << op->GetName() << ":"
<< input_map[index];
(void)op->SetInput(input_map[index], *(handle.op), handle.out);
}
return SUCCESS;
}
return NOT_FOUND;
}
Status SetNormalOpInput(const OperatorPtr &op, int index, const OutHandler &handle) {
MS_EXCEPTION_IF_NULL(op);
auto it = input_map_.find(index);
if ((handle.op != nullptr) && (it != input_map_.end())) {
if (handle.out.empty()) {
MS_LOG(DEBUG) << "Link op " << handle.op->GetName() << " to " << op->GetName() << ":" << it->second.name;
it->second.set_op(op, handle.op);
} else {
MS_LOG(DEBUG) << "Link op " << handle.op->GetName() << ":" << handle.out << " to " << op->GetName() << ":"
<< it->second.name;
it->second.set_handle(op, handle);
}
return SUCCESS;
}
return NOT_FOUND;
}
int setInput(const OperatorPtr &op, int index, const OutHandler &handle) override {
if (IsCustomOp(op)) {
auto cus_op = std::dynamic_pointer_cast<CustomOperator>(op);
return static_cast<int>(SetCustomOpInput(cus_op, index, handle));
} else {
return static_cast<int>(SetNormalOpInput(op, index, handle));
}
}
int setInput(const OperatorPtr &op, int index, const std::shared_ptr<std::vector<OutHandler>> &handler_vec) override {
MS_EXCEPTION_IF_NULL(handler_vec);
if (IsCustomOp(op)) {
MS_LOG(ERROR) << "Custom Op do not support dynamic input";
return static_cast<int>(FAILED);
}
MS_EXCEPTION_IF_NULL(op);
auto it = dyn_input_map_.find(index);
if (it != dyn_input_map_.end()) {
it->second.create_dyn_input(op, static_cast<unsigned int>(handler_vec->size()));
for (unsigned int i = 0; i < handler_vec->size(); ++i) {
OutHandler h = (*handler_vec)[i];
MS_EXCEPTION_IF_NULL(h.op);
if (h.out.empty()) {
MS_LOG(DEBUG) << "Link op " << h.op->GetName() << " to " << op->GetName() << ":" << it->second.name;
it->second.set_op(op, (i) /* index start from 0 */, h.op);
} else {
MS_LOG(DEBUG) << "Link op " << h.op->GetName() << ":" << h.out << " to " << op->GetName() << ":"
<< it->second.name;
it->second.set_handle(op, i, h);
}
}
return 0;
}
return static_cast<int>(NOT_FOUND);
}
OutHandler getOutput(const OperatorPtr &op, int index) override {
MS_EXCEPTION_IF_NULL(op);
if (IsCustomOp(op)) {
return getCustomOutput(op, index);
}
return getNormalOutput(op, index);
}
OutHandler getCustomOutput(const OperatorPtr &op, int index) {
MS_EXCEPTION_IF_NULL(op);
auto it = cus_output_map_.find(op->GetOpType());
if (it == cus_output_map_.end()) {
MS_LOG(ERROR) << "OpAdpator(" << op->GetName() << ") has both OUTPUT is not supported!";
return OutHandler();
}
std::unordered_map<int, std::string> &output_map = it->second;
if ((output_map.find(index) != output_map.end())) {
return OutHandler(op, output_map[index]);
}
MS_LOG(ERROR) << "OpAdpator(" << op->GetName() << ") has no OUTPUT index(" << index << ")!";
return OutHandler();
}
OutHandler getNormalOutput(const OperatorPtr &op, int index) {
MS_EXCEPTION_IF_NULL(op);
if (!dyn_output_map_.empty() && !output_map_.empty()) {
MS_LOG(ERROR) << "OpAdpator(" << op->GetName() << ") has both OUTPUT and DYN_OUTPUT is not supported!";
return OutHandler();
}
auto it = output_map_.find(index);
if (it != output_map_.end()) {
return OutHandler(op, it->second.name);
} else if (!dyn_output_map_.empty()) {
return OutHandler(op, dyn_output_map_.begin()->second.name + std::to_string(index));
} else {
MS_LOG(ERROR) << "OpAdpator(" << op->GetName() << ") has no OUTPUT and DYN_OUTPUT index(" << index << ")!";
return OutHandler();
}
}
Status UpdateSingleOutputDesc(const OperatorPtr &op, const abstract::BaseShapePtr &shp, const TypePtr &type) {
MS_EXCEPTION_IF_NULL(type);
std::string format = "NCHW";
if (op->GetOpType() == kExtractImagePatchesOpName) {
format = "NHWC";
}
auto desc = CreateOutputDesc(dyn_cast<abstract::Shape>(shp), type, format);
if (desc == nullptr) {
MS_LOG(ERROR) << "Update output descriptor failed!";
return FAILED;
}
if (IsCustomOp(op)) {
if (cus_output_map_.find(op->GetOpType()) == cus_output_map_.end() ||
(cus_output_map_[op->GetOpType()].empty())) {
MS_LOG(ERROR) << "This op does not create custom output map";
return FAILED;
}
auto cus_op = std::dynamic_pointer_cast<CustomOperator>(op);
MS_EXCEPTION_IF_NULL(cus_op);
std::unordered_map<int, std::string> output_map = cus_output_map_[op->GetOpType()];
(void)cus_op->UpdateOutputDesc(output_map[0], *desc);
} else {
if (output_map_.empty()) {
MS_LOG(INFO) << "This op does not have output map";
return FAILED;
}
output_map_.begin()->second.update_out_desc(op, *desc);
}
return SUCCESS;
}
size_t GetCustomOpOutputSize(const CusOperatorPtr &cus_op) {
MS_EXCEPTION_IF_NULL(cus_op);
if (cus_output_map_.find(cus_op->GetOpType()) == cus_output_map_.end()) {
MS_LOG(ERROR) << "This op does not create custom output map";
return 0;
}
size_t output_size = cus_output_map_[cus_op->GetOpType()].size();
return output_size;
}
std::shared_ptr<GeTensorDesc> CreateOutputDesc(const abstract::ShapePtr &shape_ptr, const TypePtr &type,
const std::string &format) {
if (shape_ptr == nullptr) {
MS_LOG(ERROR) << "Shape ptr is nullptr";
return nullptr;
}
if (type == nullptr) {
MS_LOG(ERROR) << "Type ptr is nullptr";
return nullptr;
}
TypeId me_type = type->type_id();
if (kObjectTypeTensorType == me_type) {
me_type = dyn_cast<TensorType>(type)->element()->type_id();
}
auto desc = TransformUtil::GetGeTensorDesc(shape_ptr->shape(), me_type, format);
return desc;
}
Status UpdateMultiOutputDesc(const OperatorPtr &op, const abstract::BaseShapePtr &shp, const TypePtr &type) {
auto tuple_shp = dyn_cast<abstract::TupleShape>(shp);
MS_EXCEPTION_IF_NULL(tuple_shp);
size_t output_size = 0;
bool is_custom_op = IsCustomOp(op);
if (is_custom_op) {
output_size = GetCustomOpOutputSize(std::dynamic_pointer_cast<CustomOperator>(op));
} else {
output_size = output_map_.size();
}
if (output_size == 0) {
MS_LOG(INFO) << "This op does not have output map";
return FAILED;
}
if (output_size != tuple_shp->shape().size()) {
MS_LOG(ERROR) << "output_map is not equal tuple_shape size";
return FAILED;
}
std::string format = "NCHW";
if (op->GetOpType() == kTopKOpName) {
format = "NHWC";
}
for (size_t i = 0; i < tuple_shp->shape().size(); ++i) {
auto tuple_type = dyn_cast<Tuple>(type);
MS_EXCEPTION_IF_NULL(tuple_type);
TypePtr type_elem = tuple_type->elements()[i];
auto desc = CreateOutputDesc(dyn_cast<abstract::Shape>(tuple_shp->shape()[i]), type_elem, format);
if (desc == nullptr) {
MS_LOG(ERROR) << "Create output descriptor failed!";
return FAILED;
}
if (is_custom_op) {
(void)std::dynamic_pointer_cast<CustomOperator>(op)->UpdateOutputDesc(cus_output_map_[op->GetOpType()][i],
*desc);
} else {
auto it = output_map_.find(i);
if (it != output_map_.end()) {
it->second.update_out_desc(op, *desc);
}
}
}
return SUCCESS;
}
std::shared_ptr<GeTensorDesc> CreateNodeDesc(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
TypeId me_type = node->Type()->type_id();
if (kObjectTypeTensorType == me_type) {
me_type = dyn_cast<TensorType>(node->Type())->element()->type_id();
}
if (me_type <= kNumberTypeBegin || me_type >= kNumberTypeEnd) {
return nullptr;
}
std::vector<int> shape;
auto shape_ptr = dyn_cast<abstract::Shape>(node->Shape());
if (nullptr != shape_ptr) {
shape = shape_ptr->shape();
}
auto desc = TransformUtil::GetGeTensorDesc(shape, me_type, "NCHW");
if (desc == nullptr) {
MS_LOG(ERROR) << "Update output descriptor failed!";
return nullptr;
}
return desc;
}
void UpdateNormalOpInputDesc(const OperatorPtr &op, const AnfNodePtr node) {
if (op == nullptr) {
MS_LOG(ERROR) << "op is nullptr";
return;
}
MS_EXCEPTION_IF_NULL(node);
auto inputs = node->cast<CNodePtr>()->inputs();
for (size_t i = 1; i < inputs.size(); ++i) {
auto it = input_map_.find(i);
if (it != input_map_.end()) {
auto desc = CreateNodeDesc(inputs[i]);
if (desc == nullptr) {
continue;
}
if (op->GetOpType() == kExtractImagePatchesOpName) {
desc->SetFormat(ge::Format::FORMAT_NHWC);
}
it->second.update_input_desc(op, *desc);
}
}
}
void UpdateCustomOpInputDesc(const CusOperatorPtr &op, const AnfNodePtr &node) {
if (op == nullptr) {
MS_LOG(ERROR) << "op is nullptr";
return;
}
MS_EXCEPTION_IF_NULL(node);
if (cus_input_map_.find(op->GetOpType()) == cus_input_map_.end() || (cus_input_map_[op->GetOpType()].empty())) {
MS_LOG(ERROR) << "This op does not create custom input map";
return;
}
std::unordered_map<int, std::string> &input_map = cus_input_map_[op->GetOpType()];
auto inputs = node->cast<CNodePtr>()->inputs();
for (size_t i = 1; i < inputs.size(); ++i) {
if (input_map.find(i) != input_map.end()) {
auto desc = CreateNodeDesc(inputs[i]);
if (desc == nullptr) {
continue;
}
(void)op->UpdateInputDesc(input_map[i], *desc);
}
}
}
void updateInputDesc(const OperatorPtr &op, const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(op);
MS_EXCEPTION_IF_NULL(node);
if (IsCustomOp(op)) {
auto cus_op = std::dynamic_pointer_cast<CustomOperator>(op);
UpdateCustomOpInputDesc(cus_op, node);
} else {
UpdateNormalOpInputDesc(op, node);
}
}
void updateOutputDesc(const OperatorPtr &op, const abstract::BaseShapePtr &shp, const TypePtr &type,
const AnfNodePtr &node) override {
if (op == nullptr) {
MS_LOG(ERROR) << "op is nullptr";
return;
}
MS_EXCEPTION_IF_NULL(node);
MS_LOG(INFO) << "Op name is " << op->GetName();
auto normal_shape_ptr = dyn_cast<abstract::Shape>(shp);
auto no_shape_ptr = dyn_cast<abstract::NoShape>(shp);
if ((nullptr != normal_shape_ptr) || (nullptr != no_shape_ptr)) {
if (UpdateSingleOutputDesc(op, shp, type) != SUCCESS) {
return;
}
} else if (nullptr != dyn_cast<abstract::TupleShape>(shp)) {
if (UpdateMultiOutputDesc(op, shp, type) != SUCCESS) {
return;
}
} else {
MS_LOG(WARNING) << "Update output desc failed, unknow output shape type";
return;
}
MS_EXCEPTION_IF_NULL(node);
if (!node->isa<CNode>()) {
return;
}
// Need to update input_desc while the output_desc is updated
updateInputDesc(op, node);
}
int setAttr(const OperatorPtr &op, const std::string &attrKey, const ValuePtr &attrValue) override {
auto it = attr_map_.find(attrKey);
if (it != attr_map_.end()) {
// switch case for each avalilable attribute type
MS_LOG(INFO) << "Set attr: " << attrKey << "(" << it->second.name << "), value: " << attrValue->ToString();
AddAttrToDrawGraph(attrKey + std::string("=") + attrValue->ToString());
it->second.set_attr(op, attrValue);
return 0;
}
return static_cast<int>(NOT_FOUND);
}
int SetCustomOpAttr(const CusOperatorPtr &op, const PrimitivePtr &prim) {
enum ValueType {
SINGLE_VALUE = 0,
SEQUEUE_VALUE,
UNKNOWN_VALUE,
};
MS_EXCEPTION_IF_NULL(prim);
MS_EXCEPTION_IF_NULL(op);
ValueType value_type = SINGLE_VALUE;
for (auto item : prim->attrs()) {
if (item.second->isa<Int32Imm>()) {
(void)op->SetAttr(item.first, GetValue<int>(item.second));
} else if (item.second->isa<StringImm>()) {
(void)op->SetAttr(item.first, GetValue<std::string>(item.second));
} else if (item.second->isa<BoolImm>()) {
(void)op->SetAttr(item.first, GetValue<bool>(item.second));
} else if (item.second->isa<FP32Imm>()) {
(void)op->SetAttr(item.first, GetValue<float>(item.second));
} else if (item.second->isa<ValueSequeue>()) {
value_type = SEQUEUE_VALUE;
auto val_seq = item.second->cast<ValueSequeuePtr>();
if ((*val_seq)[0]->isa<StringImm>()) {
(void)op->SetAttr(item.first, GetValue<const std::vector<std::string>>(item.second));
} else if ((*val_seq)[0]->isa<FP32Imm>()) {
(void)op->SetAttr(item.first, GetValue<const std::vector<float>>(item.second));
} else if ((*val_seq)[0]->isa<Int32Imm>()) {
(void)op->SetAttr(item.first, GetValue<const std::vector<int>>(item.second));
} else if ((*val_seq)[0]->isa<BoolImm>()) {
(void)op->SetAttr(item.first, GetValue<const std::vector<bool>>(item.second));
} else {
MS_LOG(EXCEPTION) << "Unsupported custom attribute type in adaptor, prim name: " << prim->name()
<< ", attr name: " << item.first << ", value: " << item.second->ToString();
}
} else {
value_type = UNKNOWN_VALUE;
MS_LOG(WARNING) << "Unsupported custom attribute type in adaptor, prim name: " << prim->name()
<< ", attr name: " << item.first << ", value: " << item.second->ToString();
return static_cast<int>(NOT_FOUND);
}
if (value_type == SINGLE_VALUE) {
AddAttrToDrawGraph(item.first + std::string("=") + item.second->ToString());
} else if (value_type == SEQUEUE_VALUE) {
AddAttrToDrawGraph(item.first + std::string("=") + "[...]");
}
}
return 0;
}
int SetNormalOpAttr(const OperatorPtr &op, const PrimitivePtr &prim) {
int ret = 0;
MS_EXCEPTION_IF_NULL(prim);
MS_EXCEPTION_IF_NULL(op);
for (auto &it : attr_map_) {
auto value = prim->GetAttr(it.first);
if (value != nullptr) {
// set attr from primitive
ret = setAttr(op, it.first, value);
if (ret) {
return ret;
}
} else {
// set attr from extra_attr
auto it_extra = extra_attr_.find(it.first);
if (it_extra != extra_attr_.end()) {
ret = setAttr(op, it.first, it_extra->second);
if (ret) {
return ret;
}
}
}
}
return 0;
}
int setAttr(const OperatorPtr &op, const PrimitivePtr &prim) override {
int ret = 0;
if (IsCustomPrim(prim)) {
auto cus_op = std::dynamic_pointer_cast<CustomOperator>(op);
ret = SetCustomOpAttr(cus_op, prim);
} else {
ret = SetNormalOpAttr(op, prim);
}
return ret;
}
int setAttr(const OperatorPtr &op, const AnfNodePtr &node) override {
// no attribute for lonely node
MS_EXCEPTION_IF_NULL(node);
if (!node->isa<CNode>()) {
return 0;
}
auto cnode = node->cast<CNodePtr>();
if (cnode == nullptr) {
return 0;
}
auto &inputs = cnode->inputs();
if (inputs.empty()) {
return 0;
}
// get Attr T from abstract of anfnode first,
// if attr "T" appears in primitive, the primitive T will cover this one
if (attr_map_.find("T") != attr_map_.end()) {
// get dtype from inputs[1], if the node has no inputs, set the attr T with output dtype
TypePtr type;
if (inputs.size() > 1) {
type = inputs[1]->Type();
} else {
type = node->Type();
}
if (type != nullptr) {
(void)setAttr(op, "T", MakeValue(type));
}
}
// set attr from primitive and ExtraAttr
if (IsValueNode<Primitive>(inputs[0])) {
// set attr from primitive
PrimitivePtr prim = GetValueNode<PrimitivePtr>(inputs[0]);
int ret = setAttr(op, prim);
if (ret != 0) {
return ret;
}
}
// set attr from const input
for (auto &it : input_attr_map_) {
if (inputs.size() <= it.first || !inputs[it.first]->isa<ValueNode>()) {
continue;
}
auto const_value = GetValueNode(inputs[it.first]);
MS_LOG(INFO) << "Set attr: input_" << it.first << "(" << it.second.name
<< "), value: " << const_value->ToString();
if (const_value->isa<None>()) {
continue;
}
AddAttrToDrawGraph(it.second.name + std::string("=") + const_value->ToString());
it.second.set_attr(op, const_value);
}
return 0;
}
std::unordered_map<std::string, ValuePtr> GetExtraAttr() override { return extra_attr_; }
private:
template <typename S>
static S ConvertAny(const ValuePtr &value, const AnyTraits<S> &) {
return GetValue<S>(value);
}
// specialization for reverse bool
static bool ConvertAny(const ValuePtr &value, const AnyTraits<bool> &, bool reverse) {
return reverse != GetValue<bool>(value);
}
template <typename P, typename Q>
static Q ConvertAny(const ValuePtr &value, const AnyTraits<P> &traits_from, const AnyTraits<Q> &traits_to) {
return ConvertAnyUtil(value, traits_from, traits_to);
}
// specialization for tensor
static GeTensor ConvertAny(const ValuePtr &value, const AnyTraits<mindspore::tensor::Tensor> &traits) {
// To-DO the format may read from ME tensor
return ConvertAnyUtil(value, traits);
}
// specialization for int
static int64_t ConvertAny(const ValuePtr &value, const AnyTraits<int64_t>) {
return static_cast<int64_t>(GetValue<int>(value));
}
// specialization for int to Vector
static std::vector<int64_t> ConvertAny(const ValuePtr &value, const std::string &name,
const AnyTraits<std::vector<int64_t>> anyTraitsInt) {
return ConvertAnyUtil(value, name, anyTraitsInt);
}
static std::vector<std::vector<int64_t>> ConvertAny(const ValuePtr &value,
const AnyTraits<std::vector<std::vector<int64_t>>>) {
MS_EXCEPTION_IF_NULL(value);
MS_LOG(INFO) << "Value: " << value->type_name();
std::vector<std::vector<int64_t>> list;
if (!value->isa<ValueTuple>()) {
MS_LOG(EXCEPTION) << "Value should be ValueTuple, but got " << value->type_name();
}
auto vec = value->cast<ValueTuplePtr>();
MS_EXCEPTION_IF_NULL(vec);
for (auto &it : vec->value()) {
MS_EXCEPTION_IF_NULL(it);
if (!it->isa<ValueTuple>()) {
MS_LOG(EXCEPTION) << "It should be ValueTuple, but got " << it->type_name();
}
auto sub_vector = it->cast<ValueTuplePtr>();
std::vector<int64_t> sublist;
for (auto &item : sub_vector->value()) {
sublist.push_back(static_cast<int64_t>(GetValue<int>(item)));
}
list.push_back(sublist);
}
return list;
}
static std::vector<int64_t> ConvertAny(const ValuePtr &value, const AnyTraits<std::vector<std::vector<int64_t>>>,
const AnyTraits<std::vector<int64_t>>) {
MS_EXCEPTION_IF_NULL(value);
MS_LOG(DEBUG) << "Value: " << value->type_name();
if (!value->isa<ValueList>()) {
MS_LOG(EXCEPTION) << "Value should be ValueList, but got " << value->type_name();
}
auto vec = value->cast<ValueListPtr>();
std::vector<int64_t> list;
for (auto &it : vec->value()) {
MS_EXCEPTION_IF_NULL(it);
if (!it->isa<ValueList>()) {
MS_LOG(EXCEPTION) << "It should be ValueList, but got " << it->type_name();
}
auto sub_vector = it->cast<ValueListPtr>();
for (auto &item : sub_vector->value()) {
list.push_back(static_cast<int64_t>(GetValue<int>(item)));
}
}
return list;
}
static std::vector<int64_t> ConvertAny(const ValuePtr &value, const AnyTraits<std::vector<int64_t>>,
const AnyTraits<std::vector<int64_t>>) {
MS_EXCEPTION_IF_NULL(value);
MS_LOG(INFO) << "Value: " << value->type_name();
std::vector<int64_t> list;
if (value->isa<ValueSequeue>()) {
auto vec = value->cast<ValueSequeuePtr>();
MS_EXCEPTION_IF_NULL(vec);
for (auto &it : vec->value()) {
list.push_back(static_cast<int64_t>(GetValue<int>(it)));
}
return list;
}
if (value->isa<Scalar>()) {
list.push_back(static_cast<int64_t>(GetValue<int>(value)));
return list;
}
MS_LOG(EXCEPTION) << "Value should be ValueTuple or Scalar, but got " << value->type_name();
}
static std::string ConvertAny(const ValuePtr &value, const AnyTraits<std::vector<int64_t>> anyTraitsVec,
const AnyTraits<std::string> anyTraitsStr) {
return ConvertAnyUtil(value, anyTraitsVec, anyTraitsStr);
}
static std::vector<float> ConvertAny(const ValuePtr &value, const AnyTraits<std::vector<float>> anyTraitsVec,
const AnyTraits<float> anyTraitsFlo) {
return ConvertAnyUtil(value, anyTraitsVec, anyTraitsFlo);
}
static std::vector<int64_t> ConvertAny(const ValuePtr &value, const std::string &format,
const AnyTraits<std::vector<int64_t>> anyTraitsVec,
const AnyTraits<int64_t> anyTraitsInt) {
return ConvertAnyUtil(value, format, anyTraitsVec, anyTraitsInt);
}
// convert value list for value tuple to vector
template <typename P, typename Q>
static std::vector<Q> ConvertAny(const ValuePtr &value, const AnyTraits<P> &anyTraitsP,
const AnyTraits<std::vector<Q>> anyTraitsQ) {
return ConvertAnyUtil(value, anyTraitsP, anyTraitsQ);
}
static int64_t ConvertAny(const ValuePtr &value, const AnyTraits<GeEnum>) {
auto name = GetValue<std::string>(value);
auto it = enum_map_.find(name);
int v = 0;
if (it != enum_map_.end()) {
v = it->second;
}
return v;
}
static GeDataType ConvertAny(const ValuePtr &value, const AnyTraits<GEType> anyTraitsGE) {
return ConvertAnyUtil(value, anyTraitsGE);
}
// convert any value to tensor
static GeTensor ConvertAny(const ValuePtr &value, const AnyTraits<AnyValue> anyTraitsValue) {
return ConvertAnyUtil(value, anyTraitsValue);
}
static const std::unordered_map<int, InputDesc> input_map_;
static const std::unordered_map<int, DynInputDesc> dyn_input_map_;
static const std::unordered_map<int, OutputDesc> output_map_;
static const std::unordered_map<int, DynOutputDesc> dyn_output_map_;
static const std::unordered_map<std::string, AttrDesc> attr_map_;
static const std::unordered_map<std::string, int> enum_map_;
// convert input from anf graph to Attr in Operators
static const std::unordered_map<unsigned int, AttrDesc> input_attr_map_;
static std::unordered_map<std::string, std::unordered_map<int, std::string>> cus_input_map_;
static std::unordered_map<std::string, std::unordered_map<int, std::string>> cus_output_map_;
std::unordered_map<std::string, ValuePtr> extra_attr_;
std::unordered_map<std::string, int> name_counts_;
};
template <typename T>
const std::unordered_map<int, InputDesc> OpAdapter<T>::input_map_;
template <typename T>
const std::unordered_map<int, DynInputDesc> OpAdapter<T>::dyn_input_map_;
template <typename T>
const std::unordered_map<int, OutputDesc> OpAdapter<T>::output_map_;
template <typename T>
const std::unordered_map<int, DynOutputDesc> OpAdapter<T>::dyn_output_map_;
template <typename T>
const std::unordered_map<std::string, AttrDesc> OpAdapter<T>::attr_map_;
template <typename T>
const std::unordered_map<std::string, int> OpAdapter<T>::enum_map_;
template <typename T>
const std::unordered_map<unsigned int, AttrDesc> OpAdapter<T>::input_attr_map_;
template <typename T>
std::unordered_map<std::string, std::unordered_map<int, std::string>> OpAdapter<T>::cus_input_map_;
template <typename T>
std::unordered_map<std::string, std::unordered_map<int, std::string>> OpAdapter<T>::cus_output_map_;
// specialization for method
} // namespace transform
} // namespace mindspore
#endif // TRANSFORM_OP_ADAPTER_H_