Merge pull request BVLC#5075 from tsocha/master

Fix mkl issue BVLC#4836
W-Jilly · Jan 14, 2017 · 0a91794 · 0a91794
2 parents 8314363 + 2fac0d6
commit 0a91794
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/include/caffe/syncedmem.hpp b/include/caffe/syncedmem.hpp
@@ -3,6 +3,10 @@
 
 #include <cstdlib>
 
+#ifdef USE_MKL
+  #include "mkl.h"
+#endif
+
 #include "caffe/common.hpp"
 
 namespace caffe {
@@ -20,7 +24,11 @@ inline void CaffeMallocHost(void** ptr, size_t size, bool* use_cuda) {
     return;
   }
 #endif
+#ifdef USE_MKL
+  *ptr = mkl_malloc(size ? size:1, 64);
+#else
   *ptr = malloc(size);
+#endif
   *use_cuda = false;
   CHECK(*ptr) << "host allocation of size " << size << " failed";
 }
@@ -32,7 +40,11 @@ inline void CaffeFreeHost(void* ptr, bool use_cuda) {
     return;
   }
 #endif
+#ifdef USE_MKL
+  mkl_free(ptr);
+#else
   free(ptr);
+#endif
 }