1 #ifndef CAFFE2_UTILS_MKL_OPERATOR_H_ 2 #define CAFFE2_UTILS_MKL_OPERATOR_H_ 4 #include "caffe2/core/operator.h" 5 #include "caffe2/proto/caffe2.pb.h" 6 #include "caffe2/utils/mkl/mkl_dnn_cppwrapper.h" 7 #include "caffe2/utils/mkl/mkl_memory.h" 11 CAFFE_DECLARE_REGISTRY(
16 #define REGISTER_MKL_OPERATOR_CREATOR(key, ...) \ 17 CAFFE_REGISTER_CREATOR(MKLOperatorRegistry, key, __VA_ARGS__) 18 #define REGISTER_MKL_OPERATOR(name, ...) \ 19 CAFFE_REGISTER_CLASS(MKLOperatorRegistry, name, __VA_ARGS__) 20 #define REGISTER_MKL_OPERATOR_STR(str_name, ...) \ 21 CAFFE_REGISTER_TYPED_CLASS(MKLOperatorRegistry, str_name, __VA_ARGS__) 23 #define REGISTER_MKL_OPERATOR_WITH_ENGINE(name, engine, ...) \ 24 CAFFE_REGISTER_CLASS(MKLOperatorRegistry, name##_ENGINE_##engine, __VA_ARGS__) 37 return OperatorBase::template Input<MKLMemory<T>>(idx);
40 return OperatorBase::template Output<MKLMemory<T>>(idx);
46 bool Run(
int stream_id)
final {
53 err.AppendMessage(
"Error from operator: \n" + ProtoDebugString(def()));
58 virtual bool RunOnDevice() = 0;
60 inline void ExecutePrimitive() {
61 MKLDNN_SAFE_CALL(mkl::dnnExecute<T>(primitive_, resources_));
68 vector<vector<TIndex>> input_size_cache_;
74 void* resources_[dnnResourceNumber];
78 #define USE_MKLOPERATOR_FUNCTIONS(T) \ 79 USE_OPERATOR_BASE_FUNCTIONS; \ 80 using MKLOperator<T>::Input; \ 81 using MKLOperator<T>::Output; \ 82 using MKLOperator<T>::ExecutePrimitive; \ 83 using MKLOperator<T>::primitive_; \ 84 using MKLOperator<T>::input_size_cache_; \ 85 using MKLOperator<T>::buffer_; \ 86 using MKLOperator<T>::resources_ 88 #define USE_SIMPLE_MKL_CTOR_DTOR(name, T) \ 89 name(const OperatorDef& operator_def, Workspace* ws) \ 90 : MKLOperator<T>(operator_def, ws) {} \ 95 #endif // CAFFE2_UTILS_MKL_OPERATOR_H_
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Simple registry implementation in Caffe2 that uses static variables to register object creators durin...
A wrapper around an opaque MKL internal resource that has certain layouts and convertion primitives s...