caffe2/core/init.h

#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 TORCH_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<std::pair<InitFunction, const char*>>& 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<std::pair<InitFunction, const char*> > early_init_functions_;
  vector<std::pair<InitFunction, const char*> > init_functions_;
  std::unordered_map<std::string, InitFunction> named_functions_;
  bool early_init_functions_run_yet_ = false;
  bool init_functions_run_yet_ = false;
};
}  // namespace internal

TORCH_API bool unsafeRunCaffe2InitFunction(
    const char* name,
    int* pargc = nullptr,
    char*** pargv = nullptr);

class TORCH_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
 */
TORCH_API bool GlobalInitAlreadyRun();

class TORCH_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.
 */
TORCH_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.
 */
TORCH_API bool GlobalInit();
}  // namespace caffe2
#endif  // CAFFE2_CORE_INIT_H_