From e5996241780c27439e0f8f2506e4de69a552e234 Mon Sep 17 00:00:00 2001 From: DanieeelLiu Date: Mon, 2 Dec 2024 14:41:48 +0800 Subject: [PATCH] [Fix](mlu-ops): solve protobuf version compatibility. --- test/mlu_op_gtest/pb_gtest/src/parser.cpp | 4 ++++ test/mlu_op_gtest/tools/pb2prototxt.cpp | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/test/mlu_op_gtest/pb_gtest/src/parser.cpp b/test/mlu_op_gtest/pb_gtest/src/parser.cpp index 6ac6657fb..50e4865b8 100644 --- a/test/mlu_op_gtest/pb_gtest/src/parser.cpp +++ b/test/mlu_op_gtest/pb_gtest/src/parser.cpp @@ -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); diff --git a/test/mlu_op_gtest/tools/pb2prototxt.cpp b/test/mlu_op_gtest/tools/pb2prototxt.cpp index 32f047270..76e1657be 100644 --- a/test/mlu_op_gtest/tools/pb2prototxt.cpp +++ b/test/mlu_op_gtest/tools/pb2prototxt.cpp @@ -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;