diff --git a/onnxruntime/core/framework/op_node_proto_helper.cc b/onnxruntime/core/framework/op_node_proto_helper.cc index e2f9d38b70..a171289575 100644 --- a/onnxruntime/core/framework/op_node_proto_helper.cc +++ b/onnxruntime/core/framework/op_node_proto_helper.cc @@ -8,30 +8,56 @@ #include "core/graph/op.h" #include "onnx/defs/schema.h" #include "core/framework/op_kernel.h" +#include "core/framework/tensorprotoutils.h" #include "onnx/defs/schema.h" #include "gsl/span" using namespace ONNX_NAMESPACE; using namespace ::onnxruntime::common; namespace onnxruntime { -#define ORT_DEFINE_GET_ATTR(IMPL_T, T, type) \ - template <> \ - template <> \ - Status OpNodeProtoHelper::GetAttr( \ - const std::string& name, T* value) const { \ - const AttributeProto* attr = TryGetAttribute(name); \ - if (!attr) { \ - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "No attribute with name:'", name, "'is defined."); \ - } \ - if (!attr->has_##type()) { \ - return Status(ONNXRUNTIME, FAIL, "Attibute name and type don't match"); \ - } else { \ - *value = static_cast(attr->type()); \ - return Status::OK(); \ - } \ +template +bool HasTyped(const AttributeProto*); + +template <> +inline bool HasTyped(const AttributeProto* attr) { + return utils::HasFloat(*attr); +} +template <> +inline bool HasTyped(const AttributeProto* attr) { + return utils::HasInt(*attr); +} +template <> +inline bool HasTyped(const AttributeProto* attr) { + return utils::HasString(*attr); +} + +template <> +inline bool HasTyped(const AttributeProto* attr) { + return utils::HasTensor(*attr); +} +template <> +inline bool HasTyped(const AttributeProto* attr) { + return utils::HasGraph(*attr); +} + +#define ORT_DEFINE_GET_ATTR(IMPL_T, T, type) \ + template <> \ + template <> \ + Status OpNodeProtoHelper::GetAttr( \ + const std::string& name, T* value) const { \ + const AttributeProto* attr = TryGetAttribute(name); \ + if (!attr) { \ + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "No attribute with name:'", name, "'is defined."); \ + } \ + if (!HasTyped(attr)) { \ + return Status(ONNXRUNTIME, FAIL, "Attibute name and type don't match"); \ + } else { \ + *value = static_cast(attr->type()); \ + return Status::OK(); \ + } \ } -#define ORT_DEFINE_GET_ATTRS(IMPL_T, T, list) \ +#define ORT_DEFINE_GET_ATTRS(IMPL_T, T, list) \ template <> \ template <> \ Status OpNodeProtoHelper::GetAttrs( \ @@ -54,7 +80,7 @@ namespace onnxruntime { if (!attr) { \ return Status(ONNXRUNTIME, FAIL, "No attribute with this name is defined."); \ } \ - ORT_ENFORCE(values.size() == attr->list##_size()); \ + ORT_ENFORCE(values.size() == attr->list##_size()); \ for (int i = 0; i < attr->list##_size(); ++i) { \ values[i] = static_cast(attr->list(i)); \ } \