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 | |
18 | namespace ONNX_NAMESPACE { |
19 | |
20 | #ifdef ONNX_USE_LITE_PROTO |
21 | using ::google::protobuf::MessageLite; |
22 | inline 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 |
29 | using ::google::protobuf::Message; |
30 | inline std::string ProtoDebugString(const Message& proto) { |
31 | return proto.ShortDebugString(); |
32 | } |
33 | #endif |
34 | |
35 | template <typename Proto> |
36 | bool 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 | |
51 | template <typename T> |
52 | inline std::vector<T> RetrieveValues(const AttributeProto& attr); |
53 | template <> |
54 | inline std::vector<int64_t> RetrieveValues(const AttributeProto& attr) { |
55 | return {attr.ints().begin(), attr.ints().end()}; |
56 | } |
57 | |
58 | template <> |
59 | inline std::vector<std::string> RetrieveValues(const AttributeProto& attr) { |
60 | return {attr.strings().begin(), attr.strings().end()}; |
61 | } |
62 | |
63 | template <> |
64 | inline std::vector<float> RetrieveValues(const AttributeProto& attr) { |
65 | return {attr.floats().begin(), attr.floats().end()}; |
66 | } |
67 | |
68 | } // namespace ONNX_NAMESPACE |
69 |