Skip to content

Commit

Permalink
[Fix](mlu-ops): solve protobuf version compatibility. (#1165)
Browse files Browse the repository at this point in the history
  • Loading branch information
DanieeelLiu authored Dec 2, 2024
1 parent a3e0420 commit ba823ed
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 0 deletions.
4 changes: 4 additions & 0 deletions test/mlu_op_gtest/pb_gtest/src/parser.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1072,7 +1072,11 @@ bool Parser::readMessageFromFile(const std::string &filename, Node *proto) {
google::protobuf::io::FileInputStream input(fd);
if (strEndsWith(filename, ".pb")) {
google::protobuf::io::CodedInputStream coded_input(&input);
#if GOOGLE_PROTOBUF_VERSION > 3005000
coded_input.SetTotalBytesLimit(INT_MAX);
#elif GOOGLE_PROTOBUF_VERSION
coded_input.SetTotalBytesLimit(INT_MAX, INT_MAX - 1);
#endif
status = proto->ParseFromCodedStream(&coded_input);
} else if (strEndsWith(filename, ".prototxt")) {
status = google::protobuf::TextFormat::Parse(&input, proto);
Expand Down
4 changes: 4 additions & 0 deletions test/mlu_op_gtest/tools/pb2prototxt.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -185,7 +185,11 @@ bool readIn(const std::string &filename, mluoptest::Node *proto) {
new google::protobuf::io::CodedInputStream(input);
// Total bytes hard limit / warning limit are set to 1GB and 512MB, as same
// as tensorflow.
#if GOOGLE_PROTOBUF_VERSION > 3005000
coded_stream->SetTotalBytesLimit(INT_MAX);
#elif GOOGLE_PROTOBUF_VERSION
coded_stream->SetTotalBytesLimit(INT_MAX, 512LL << 20);
#endif
proto->ParseFromCodedStream(coded_stream);
delete input;
delete coded_stream;
Expand Down

0 comments on commit ba823ed

Please sign in to comment.