fix torch.config can't respect USE_MKLDNN flag issue (#75001)

Fixes https://github.com/pytorch/pytorch/issues/74949, which reports that torch.config can't respect USE_MKLDNN flag.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/75001
Approved by: https://github.com/malfet
This commit is contained in:
zhang, xiaobing
2022-07-17 15:00:47 +00:00
committed by PyTorch MergeBot
parent 4655c3bace
commit 86b86202b5
22 changed files with 40 additions and 41 deletions

View File

@ -17,7 +17,7 @@ if(BUILD_CAFFE2_OPS)
torch_set_target_props(caffe2_detectron_ops_gpu)
target_link_libraries(caffe2_detectron_ops_gpu PRIVATE torch ${OpenMP_link})
if(CAFFE2_USE_MKLDNN)
if(USE_MKLDNN)
target_link_libraries(caffe2_detectron_ops_gpu PRIVATE caffe2::mkldnn)
endif()
install(TARGETS caffe2_detectron_ops_gpu DESTINATION lib)
@ -33,7 +33,7 @@ if(BUILD_CAFFE2_OPS)
${Detectron_HIP_SRCS})
torch_set_target_props(caffe2_detectron_ops_hip)
target_compile_options(caffe2_detectron_ops_hip PRIVATE ${HIP_CXX_FLAGS})
if(CAFFE2_USE_MKLDNN)
if(USE_MKLDNN)
target_link_libraries(caffe2_detectron_ops_hip PRIVATE caffe2::mkldnn)
endif()
target_link_libraries(caffe2_detectron_ops_hip PRIVATE torch)
@ -46,7 +46,7 @@ if(BUILD_CAFFE2_OPS)
endif()
torch_set_target_props(caffe2_detectron_ops)
target_link_libraries(caffe2_detectron_ops PRIVATE torch ${OpenMP_link})
if(CAFFE2_USE_MKLDNN)
if(USE_MKLDNN)
target_link_libraries(caffe2_detectron_ops PRIVATE caffe2::mkldnn)
endif()
install(TARGETS caffe2_detectron_ops DESTINATION lib)

View File

@ -15,13 +15,13 @@
*/
#include "upsample_nearest_op.h"
#ifdef CAFFE2_USE_MKLDNN
#ifdef USE_MKLDNN
#include "caffe2/ideep/operators/operator_fallback_ideep.h"
#include "caffe2/ideep/utils/ideep_operator.h"
#endif
namespace caffe2 {
#ifdef CAFFE2_USE_MKLDNN
#ifdef USE_MKLDNN
REGISTER_IDEEP_OPERATOR(
UpsampleNearest,
IDEEPFallbackOp<UpsampleNearestOp<float, CPUContext>>);