1/*
2 * SPDX-License-Identifier: Apache-2.0
3 */
4
5#pragma once
6
7#include <google/protobuf/io/coded_stream.h>
8#include <google/protobuf/io/zero_copy_stream_impl_lite.h>
9
10#include "onnx/onnx_pb.h"
11
12#ifdef ONNX_USE_LITE_PROTO
13#include <google/protobuf/message_lite.h>
14#else // ONNX_USE_LITE_PROTO
15#include <google/protobuf/message.h>
16#endif // !ONNX_USE_LITE_PROTO
17
18namespace ONNX_NAMESPACE {
19
20#ifdef ONNX_USE_LITE_PROTO
21using ::google::protobuf::MessageLite;
22inline std::string ProtoDebugString(const MessageLite& proto) {
23 // Since the MessageLite interface does not support reflection, there is very
24 // little information that this and similar methods can provide.
25 // But when using lite proto this is the best we can provide.
26 return proto.ShortDebugString();
27}
28#else
29using ::google::protobuf::Message;
30inline std::string ProtoDebugString(const Message& proto) {
31 return proto.ShortDebugString();
32}
33#endif
34
35template <typename Proto>
36bool ParseProtoFromBytes(Proto* proto, const char* buffer, size_t length) {
37 ::google::protobuf::io::ArrayInputStream input_stream(buffer, static_cast<int>(length));
38 ::google::protobuf::io::CodedInputStream coded_stream(&input_stream);
39 int total_bytes_limit = (2048LL << 20) - 1;
40#if GOOGLE_PROTOBUF_VERSION >= 3011000
41 // Only take one parameter since protobuf 3.11
42 coded_stream.SetTotalBytesLimit(total_bytes_limit);
43#else
44 // Total bytes hard limit / warning limit are set to 2GB and 512MB respectively.
45 coded_stream.SetTotalBytesLimit(total_bytes_limit, 512LL << 20);
46#endif
47
48 return proto->ParseFromCodedStream(&coded_stream);
49}
50
51template <typename T>
52inline std::vector<T> RetrieveValues(const AttributeProto& attr);
53template <>
54inline std::vector<int64_t> RetrieveValues(const AttributeProto& attr) {
55 return {attr.ints().begin(), attr.ints().end()};
56}
57
58template <>
59inline std::vector<std::string> RetrieveValues(const AttributeProto& attr) {
60 return {attr.strings().begin(), attr.strings().end()};
61}
62
63template <>
64inline std::vector<float> RetrieveValues(const AttributeProto& attr) {
65 return {attr.floats().begin(), attr.floats().end()};
66}
67
68} // namespace ONNX_NAMESPACE
69