Caffe2 - C++ API
A deep learning, cross platform ML framework
mkl_memory.cc
1 #include "caffe2/mkl/mkl_utils.h"
2 
3 #include "caffe2/core/init.h"
4 #include "caffe2/core/tensor.h"
5 
6 #ifdef CAFFE2_HAS_MKL_DNN
7 
8 CAFFE2_DEFINE_bool(
9  caffe2_mkl_implicit_layout_change, false,
10  "Controls the behavior when we call View() on an MKLMemory: if it is set "
11  "true, then the View() function will actually change the underlying "
12  "storage. If it is set false, an implicit copy is triggered but the "
13  "original storage is not affected."
14  );
15 
16 namespace caffe2 {
17 
18 CAFFE_KNOWN_TYPE(mkl::MKLMemory<float>);
19 CAFFE_KNOWN_TYPE(mkl::MKLMemory<double>);
20 
21 template <typename T>
22 static vector<TIndex> GetMKLTensorInfo(
23  const void* c,
24  bool* shares_data,
25  size_t* capacity,
26  DeviceOption* device) {
27  const mkl::MKLMemory<T>* tc = static_cast<const mkl::MKLMemory<T>*>(c);
28  // it's too hard to get sharing info from mkl::MKLMemory
29  *shares_data = false;
30  *capacity = tc->size() * sizeof(T);
31  device->set_device_type(MKLDNN);
32  device->set_cuda_gpu_id(0);
33  return tc->dims();
34 }
35 
36 template <typename T>
37 static TypeMeta GetMKLTensorType(const void*) {
38  return TypeMeta::Make<T>();
39 }
40 
41 template <typename T>
42 static void RegisterForType() {
43  RegisterTypeCallFunction(
44  TypeMeta::Id<mkl::MKLMemory<T>>(), GetMKLTensorType<T>);
45  RegisterTensorInfoFunction(
46  TypeMeta::Id<mkl::MKLMemory<T>>(), GetMKLTensorInfo<T>);
47 }
48 
49 static bool Caffe2InitializeMKL(int*, char***) {
50  RegisterForType<float>();
51  RegisterForType<double>();
52  return true;
53 }
54 
55 REGISTER_CAFFE2_INIT_FUNCTION(
56  InitMKLDNNContext,
57  &Caffe2InitializeMKL,
58  "Register wrappers for MKLContext");
59 
60 } // namespace caffe2
61 
62 #endif // CAFFE2_HAS_MKL_DNN
static CAFFE2_API CaffeTypeId Id()
Returns the unique id for the given type T.
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...