operator_meta.h

Go to the documentation of this file.

#pragma once

#include <functional>
#include "caffe2/core/logging.h"
#include "caffe2/core/registry.h"
#include "caffe2/proto/caffe2.pb.h"

namespace caffe2 {

using ReferenceImplementation =
    std::function<void(NetDef*, const OperatorDef&)>;

class ReferenceImplementationRegistry {
 public:
  class Register {
   public:
    Register(const std::string& name, ReferenceImplementation func) {
      CAFFE_ENFORCE(!ReferenceImplementationRegistry::getMap().count(name));
      ReferenceImplementationRegistry::getMap()[name] = func;
    }
  };

  static void Append(NetDef* net, const OperatorDef& op);
  static NetDef ConvertNet(const NetDef& net);

 private:
  ReferenceImplementationRegistry() = delete;
  static CaffeMap<std::string, ReferenceImplementation>& getMap();
};

#define TC_REFERENCE_IMPLEMENTATION(name, func)              \
  static ::caffe2::ReferenceImplementationRegistry::Register \
      CAFFE_ANONYMOUS_VARIABLE(tc_ref##name)(#name, func);
} // namespace caffe2