Caffe2 - C++ API
A deep learning, cross platform ML framework
mkl_memory.cc
1 
17 #include "caffe2/mkl/mkl_utils.h"
18 
19 #include "caffe2/core/init.h"
20 #include "caffe2/core/tensor.h"
21 
22 #ifdef CAFFE2_HAS_MKL_DNN
23 
24 CAFFE2_DEFINE_bool(
25  caffe2_mkl_implicit_layout_change, false,
26  "Controls the behavior when we call View() on an MKLMemory: if it is set "
27  "true, then the View() function will actually change the underlying "
28  "storage. If it is set false, an implicit copy is triggered but the "
29  "original storage is not affected."
30  );
31 
32 namespace caffe2 {
33 
34 CAFFE_KNOWN_TYPE(mkl::MKLMemory<float>);
35 CAFFE_KNOWN_TYPE(mkl::MKLMemory<double>);
36 
37 template <typename T>
38 static vector<TIndex> GetMKLTensorInfo(
39  const void* c,
40  bool* shares_data,
41  size_t* capacity,
42  DeviceOption* device) {
43  const mkl::MKLMemory<T>* tc = static_cast<const mkl::MKLMemory<T>*>(c);
44  // it's too hard to get sharing info from mkl::MKLMemory
45  *shares_data = false;
46  *capacity = tc->size() * sizeof(T);
47  device->set_device_type(MKLDNN);
48  device->set_cuda_gpu_id(0);
49  return tc->dims();
50 }
51 
52 template <typename T>
53 static TypeMeta GetMKLTensorType(const void*) {
54  return TypeMeta::Make<T>();
55 }
56 
57 template <typename T>
58 static void RegisterForType() {
59  RegisterTypeCallFunction(
60  TypeMeta::Id<mkl::MKLMemory<T>>(), GetMKLTensorType<T>);
61  RegisterTensorInfoFunction(
62  TypeMeta::Id<mkl::MKLMemory<T>>(), GetMKLTensorInfo<T>);
63 }
64 
65 static bool Caffe2InitializeMKL(int*, char***) {
66  RegisterForType<float>();
67  RegisterForType<double>();
68  return true;
69 }
70 
71 REGISTER_CAFFE2_INIT_FUNCTION(
72  InitMKLDNNContext,
73  &Caffe2InitializeMKL,
74  "Register wrappers for MKLContext");
75 
76 } // namespace caffe2
77 
78 #endif // CAFFE2_HAS_MKL_DNN
static CAFFE2_API CaffeTypeId Id()
Returns the unique id for the given type T.
Copyright (c) 2016-present, Facebook, Inc.