Post process onnx proto (#8064)

* Post processing onnx generated protobuf files to hide global symbols * . * .
sungin-h · Jun 2, 2018 · 5e35fbf · 5e35fbf
1 parent 01f5ee7
commit 5e35fbf
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 11 deletions.
diff --git a/caffe2/utils/proto_utils.cc b/caffe2/utils/proto_utils.cc
@@ -28,6 +28,17 @@ const ::std::string& GetEmptyStringAlreadyInited() {
 
 }  // namespace caffe
 
+namespace ONNX_NAMESPACE {
+
+// ONNX wrapper functions for protobuf's GetEmptyStringAlreadyInited() function
+// used to avoid duplicated global variable in the case when protobuf
+// is built with hidden visibility.
+const ::std::string& GetEmptyStringAlreadyInited() {
+  return ::google::protobuf::internal::GetEmptyStringAlreadyInited();
+}
+
+}  // namespace ONNX_NAMESPACE
+
 namespace caffe2 {
 
 // Caffe2 wrapper functions for protobuf's GetEmptyStringAlreadyInited() function

diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
@@ -664,6 +664,11 @@ if (CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO)
   # We will build onnx as static libs and embed it directly into the binary.
   set(BUILD_SHARED_LIBS OFF)
   set(ONNX_USE_MSVC_STATIC_RUNTIME ${CAFFE2_USE_MSVC_STATIC_RUNTIME})
+  # Do not do post-processing if caffe2 is not included in the build,
+  # otherwise the caffe2 protobuf symbols will not be found
+  if (BUILD_CAFFE2 AND CAFFE2_LINK_LOCAL_PROTOBUF)
+    set(ONNX_PROTO_POST_BUILD_SCRIPT ${PROJECT_SOURCE_DIR}/cmake/ProtoBufPatch.cmake)
+  endif()
   add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/../third_party/onnx)
   include_directories(${ONNX_INCLUDE_DIRS})
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DONNX_NAMESPACE=${ONNX_NAMESPACE}")

diff --git a/cmake/ProtoBuf.cmake b/cmake/ProtoBuf.cmake
@@ -181,7 +181,7 @@ function(caffe2_protobuf_generate_cpp_py srcs_var hdrs_var python_var)
 
         # If we remove all reference to these pb.h files from external
         # libraries and binaries this rewrite can be removed.
-        COMMAND ${CMAKE_COMMAND} -DFILENAME=${CMAKE_CURRENT_BINARY_DIR}/${fil_we}.pb.h -P ${PROJECT_SOURCE_DIR}/cmake/ProtoBufPatch.cmake
+        COMMAND ${CMAKE_COMMAND} -DFILENAME=${CMAKE_CURRENT_BINARY_DIR}/${fil_we}.pb.h -DNAMESPACES=caffe\;caffe2 -P ${PROJECT_SOURCE_DIR}/cmake/ProtoBufPatch.cmake
 
         DEPENDS ${CAFFE2_PROTOC_EXECUTABLE} ${abs_fil}
         COMMENT "Running C++/Python protocol buffer compiler on ${fil}" VERBATIM )

diff --git a/cmake/ProtoBufPatch.cmake b/cmake/ProtoBufPatch.cmake
@@ -3,20 +3,21 @@
 #   cmake -DFILENAME=caffe2.pb.h -P ProtoBufPatch.cmake
 
 file(READ ${FILENAME} content)
+
 string(
   REPLACE
   "::google::protobuf::internal::GetEmptyStringAlreadyInited"
   "GetEmptyStringAlreadyInited"
   content
   "${content}")
-string(REPLACE
-	"namespace caffe2 {"
-	"namespace caffe2 { const ::std::string& GetEmptyStringAlreadyInited(); "
-	content
-	"${content}")
-string(REPLACE
-	"namespace caffe {"
-	"namespace caffe { const ::std::string& GetEmptyStringAlreadyInited(); "
-	content
-	"${content}")
+
+foreach(ns ${NAMESPACES})
+  string(
+    REPLACE
+    "namespace ${ns} {"
+    "namespace ${ns} { const ::std::string& GetEmptyStringAlreadyInited(); "
+    content
+    "${content}")
+endforeach()
+
 file(WRITE ${FILENAME} "${content}")