From 57b51db8d78f46f4037c678cb8ea0265f1cc48e2 Mon Sep 17 00:00:00 2001
From: Yangqing Jia <jiayq84@gmail.com>
Date: Fri, 21 Apr 2017 03:19:50 -0700
Subject: [PATCH] Add a guard function to check Caffe2 linking setup.

Summary:
This helps diagnosing issues like #346
Closes https://github.com/caffe2/caffe2/pull/354

Differential Revision: D4928347

Pulled By: Yangqing

fbshipit-source-id: b45685f1da18cbc49be293260b1fc2268fe5cd4c
---
 caffe2/core/init.cc     |  2 ++
 caffe2/core/operator.cc |  2 ++
 caffe2/core/operator.h  | 26 ++++++++++++++++++++++++++
 3 files changed, 30 insertions(+)

diff --git a/caffe2/core/init.cc b/caffe2/core/init.cc
index bac6766aadc..c7c07329df1 100644
--- a/caffe2/core/init.cc
+++ b/caffe2/core/init.cc
@@ -1,4 +1,5 @@
 #include "caffe2/core/init.h"
+#include "caffe2/core/operator.h" // for StaticLinkingProtector
 
 #ifndef CAFFE2_BUILD_STRING
 #define CAFFE2_BUILD_STRING "build_version_not_set"
@@ -14,6 +15,7 @@ Caffe2InitializeRegistry* Caffe2InitializeRegistry::Registry() {
 
 bool GlobalInit(int* pargc, char*** pargv) {
   static bool global_init_was_already_run = false;
+  static StaticLinkingProtector g_protector;
   if (global_init_was_already_run) {
     VLOG(1) << "GlobalInit has already been called: did you double-call?";
     return true;
diff --git a/caffe2/core/operator.cc b/caffe2/core/operator.cc
index 5706360c301..63a93177e2b 100644
--- a/caffe2/core/operator.cc
+++ b/caffe2/core/operator.cc
@@ -54,10 +54,12 @@ unique_ptr<OperatorBase> TryCreateOperator(
     return nullptr;
   }
 }
+
 }  // namespace
 
 unique_ptr<OperatorBase> CreateOperator(
     const OperatorDef& operator_def, Workspace* ws) {
+  static StaticLinkingProtector g_protector;
   // first, check with OpSchema if the operator is legal.
   auto* schema = OpSchemaRegistry::Schema(operator_def.type());
   if (schema) {
diff --git a/caffe2/core/operator.h b/caffe2/core/operator.h
index faa957f3bd3..47b4ab3d60e 100644
--- a/caffe2/core/operator.h
+++ b/caffe2/core/operator.h
@@ -449,6 +449,32 @@ CAFFE_DECLARE_REGISTRY(
 #define REGISTER_CUDNN_OPERATOR(name, ...) \
   REGISTER_CUDA_OPERATOR_WITH_ENGINE(name, CUDNN, __VA_ARGS__)
 
+// StaticLinkingProtector is a helper class that ensures that the Caffe2
+// library is linked correctly with whole archives (in the case of static
+// linking). What happens is that when CreateOperator is called for the first
+// time, it instantiates an OperatorLinkingProtector object to check if the
+// operator registry is empty. If it is empty, this means that we are not
+// properly linking the library.
+//
+// You should not need to use this class.
+struct StaticLinkingProtector {
+  StaticLinkingProtector() {
+    const int registered_ops = CPUOperatorRegistry()->Keys().size();
+    // Note: this is a check failure instead of an exception, because if
+    // the linking is wrong, Caffe2 won't be able to run properly anyway,
+    // so it's better to fail loud.
+    // If Caffe2 is properly linked with whole archive, there should be more
+    // than zero registered ops.
+    if (registered_ops == 0) {
+      LOG(FATAL) <<
+        "You might have made a build error: the Caffe2 library does not seem "
+        "to be linked with whole-static library option. To do so, use "
+        "-Wl,-force_load (clang) or -Wl,--whole-archive (gcc) to link the "
+        "Caffe2 library.";
+    }
+  }
+};
+
 // An exception that can be thrown by an operator constructor that notifies
 // that it does not support the given setting. This can be usually used for
 // specific engines that only implement a subset of the features required by