#ifndef CAFFE2_CORE_INIT_H_ #define CAFFE2_CORE_INIT_H_ #include "caffe2/core/common.h" #include "caffe2/core/flags.h" #include "caffe2/core/logging.h" namespace caffe2 { namespace internal { class CAFFE2_API Caffe2InitializeRegistry { public: typedef bool (*InitFunction)(int*, char***); // Registry() is defined in .cpp file to make registration work across // multiple shared libraries loaded with RTLD_LOCAL static Caffe2InitializeRegistry* Registry(); void Register( InitFunction function, bool run_early, const char* description, const char* name = nullptr) { if (name) { named_functions_[name] = function; } if (run_early) { // Disallow registration after GlobalInit of early init functions CAFFE_ENFORCE(!early_init_functions_run_yet_); early_init_functions_.emplace_back(function, description); } else { if (init_functions_run_yet_) { // Run immediately, since GlobalInit already ran. This should be // rare but we want to allow it in some cases. LOG(WARNING) << "Running init function after GlobalInit: " << description; // TODO(orionr): Consider removing argc and argv for non-early // registration. Unfortunately that would require a new InitFunction // typedef, so not making the change right now. // // Note that init doesn't receive argc and argv, so the function // might fail and we want to raise an error in that case. int argc = 0; char** argv = nullptr; bool success = (function)(&argc, &argv); CAFFE_ENFORCE(success); } else { // Wait until GlobalInit to run init_functions_.emplace_back(function, description); } } } bool RunRegisteredEarlyInitFunctions(int* pargc, char*** pargv) { CAFFE_ENFORCE(!early_init_functions_run_yet_); early_init_functions_run_yet_ = true; return RunRegisteredInitFunctionsInternal( early_init_functions_, pargc, pargv); } bool RunRegisteredInitFunctions(int* pargc, char*** pargv) { CAFFE_ENFORCE(!init_functions_run_yet_); init_functions_run_yet_ = true; return RunRegisteredInitFunctionsInternal(init_functions_, pargc, pargv); } bool RunNamedFunction(const char* name, int* pargc, char*** pargv) { if (named_functions_.count(name)) { return named_functions_[name](pargc, pargv); } return false; } private: // Run all registered initialization functions. This has to be called AFTER // all static initialization are finished and main() has started, since we are // using logging. bool RunRegisteredInitFunctionsInternal( vector>& functions, int* pargc, char*** pargv) { for (const auto& init_pair : functions) { VLOG(1) << "Running init function: " << init_pair.second; if (!(*init_pair.first)(pargc, pargv)) { LOG(ERROR) << "Initialization function failed."; return false; } } return true; } Caffe2InitializeRegistry() {} vector > early_init_functions_; vector > init_functions_; std::unordered_map named_functions_; bool early_init_functions_run_yet_ = false; bool init_functions_run_yet_ = false; }; } // namespace internal CAFFE2_API bool unsafeRunCaffe2InitFunction( const char* name, int* pargc = nullptr, char*** pargv = nullptr); class CAFFE2_API InitRegisterer { public: InitRegisterer( internal::Caffe2InitializeRegistry::InitFunction function, bool run_early, const char* description, const char* name = nullptr) { internal::Caffe2InitializeRegistry::Registry()->Register( function, run_early, description, name); } }; #define REGISTER_CAFFE2_INIT_FUNCTION(name, function, description) \ namespace { \ ::caffe2::InitRegisterer \ g_caffe2_initregisterer_##name(function, false, description, #name); \ } // namespace #define REGISTER_CAFFE2_EARLY_INIT_FUNCTION(name, function, description) \ namespace { \ ::caffe2::InitRegisterer \ g_caffe2_initregisterer_##name(function, true, description, #name); \ } // namespace /** * @brief Determine whether GlobalInit has already been run */ CAFFE2_API bool GlobalInitAlreadyRun(); class CAFFE2_API GlobalInitIsCalledGuard { public: GlobalInitIsCalledGuard() { if (!GlobalInitAlreadyRun()) { LOG(WARNING) << "Caffe2 GlobalInit should be run before any other API calls."; } } }; /** * @brief Initialize the global environment of caffe2. * * Caffe2 uses a registration pattern for initialization functions. Custom * initialization functions should take the signature * bool (*func)(int*, char***) * where the pointers to argc and argv are passed in. Caffe2 then runs the * initialization in three phases: * (1) Functions registered with REGISTER_CAFFE2_EARLY_INIT_FUNCTION. Note that * since it is possible the logger is not initialized yet, any logging in * such early init functions may not be printed correctly. * (2) Parses Caffe-specific commandline flags, and initializes caffe logging. * (3) Functions registered with REGISTER_CAFFE2_INIT_FUNCTION. * If there is something wrong at each stage, the function returns false. If * the global initialization has already been run, the function returns false * as well. * * GlobalInit is re-entrant safe; a re-entrant call will no-op and exit. * * GlobalInit is safe to call multiple times but not idempotent; * successive calls will parse flags and re-set caffe2 logging levels from * flags as needed, but NOT re-run early init and init functions. * * GlobalInit is also thread-safe and can be called concurrently. */ CAFFE2_API bool GlobalInit(int* pargc, char*** argv); /** * @brief Initialize the global environment without command line arguments * * This is a version of the GlobalInit where no argument is passed in. * On mobile devices, use this global init, since we cannot pass the * command line options to caffe2, no arguments are passed. */ CAFFE2_API bool GlobalInit(); } // namespace caffe2 #endif // CAFFE2_CORE_INIT_H_