1 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS 2 #include <ATen/Config.h> 3 #include <c10/core/Allocator.h> 4 5 #if AT_MKLDNN_ENABLED() 6 7 // needs to be included only once in library. 8 #include <ideep_pin_singletons.hpp> 9 10 using namespace ideep; 11 12 RegisterEngineAllocator cpu_alloc( 13 engine::cpu_engine(), __anon9d6929e00102(size_t size) 14 [](size_t size) { 15 return c10::GetAllocator(c10::DeviceType::CPU)->raw_allocate(size); 16 }, __anon9d6929e00202(void* p) 17 [](void* p) { 18 c10::GetAllocator(c10::DeviceType::CPU)->raw_deallocate(p); 19 } 20 ); 21 22 namespace at::native::mkldnn{ 23 void clear_computation_cache(); 24 clear_computation_cache()25void clear_computation_cache() { 26 // Reset computation_cache for forward convolutions 27 // As it also caches max number of OpenMP workers 28 ideep::convolution_forward::t_store().clear(); 29 } 30 31 } // namespace at::native::mkldnn 32 33 #endif // AT_MKLDNN_ENABLED() 34