From a579e523a940d60bb2b5a3ee3f660ccf19fbf192 Mon Sep 17 00:00:00 2001 From: Jacek Czaja Date: Mon, 3 Oct 2022 09:45:18 +0200 Subject: [PATCH] Requantize to use Memory Desc in Tensors (#46608) * - some more MD changes * - lint * - compilation fixes * - compilation fixes * - lint * - fix --- paddle/fluid/framework/data_layout_transform.cc | 2 -- paddle/fluid/framework/tensor_util.cc | 4 +++- paddle/fluid/operators/mkldnn/requantize_mkldnn_op.cc | 3 +-- paddle/fluid/operators/transfer_layout_op.h | 9 +++++++-- paddle/phi/kernels/funcs/data_layout_transform.cc | 2 -- 5 files changed, 11 insertions(+), 9 deletions(-) diff --git a/paddle/fluid/framework/data_layout_transform.cc b/paddle/fluid/framework/data_layout_transform.cc index 254a7abd66..1594272fc5 100644 --- a/paddle/fluid/framework/data_layout_transform.cc +++ b/paddle/fluid/framework/data_layout_transform.cc @@ -202,8 +202,6 @@ void innerTransDataLayoutFromMKLDNN(DataLayout in_layout, platform::MatchShapeToLayout(out, in_layout, out_layout); out->set_layout(DataLayout::kNCHW); - // reset format since the out tensor will be feed to non-MKLDNN OPkernel - out->set_format(MKLDNNMemoryFormat::undef); } #endif diff --git a/paddle/fluid/framework/tensor_util.cc b/paddle/fluid/framework/tensor_util.cc index efc7f685bc..0d8ed3c4eb 100644 --- a/paddle/fluid/framework/tensor_util.cc +++ b/paddle/fluid/framework/tensor_util.cc @@ -471,7 +471,9 @@ void TensorCopySync(const phi::DenseTensor& src, dst->Resize(src.dims()); dst->set_layout(src.layout()); #ifdef PADDLE_WITH_MKLDNN - dst->set_format(src.format()); + if (src.layout() == DataLayout::kMKLDNN) { + dst->set_mem_desc(src.mem_desc()); + } #endif auto src_place = src.place(); auto src_ptr = src.data(); diff --git a/paddle/fluid/operators/mkldnn/requantize_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/requantize_mkldnn_op.cc index abfef00ae1..23409db02b 100644 --- a/paddle/fluid/operators/mkldnn/requantize_mkldnn_op.cc +++ b/paddle/fluid/operators/mkldnn/requantize_mkldnn_op.cc @@ -146,8 +146,7 @@ class ReQuantOpKernel : public framework::OpKernel { reorder_p->execute(astream, *src_memory, *dst_memory); astream.wait(); - output->set_layout(framework::DataLayout::kMKLDNN); - output->set_format(platform::GetMKLDNNFormat(*dst_memory)); + output->set_mem_desc(dst_memory->get_desc()); } }; diff --git a/paddle/fluid/operators/transfer_layout_op.h b/paddle/fluid/operators/transfer_layout_op.h index a4c7b482ff..aee8592842 100644 --- a/paddle/fluid/operators/transfer_layout_op.h +++ b/paddle/fluid/operators/transfer_layout_op.h @@ -93,8 +93,13 @@ class TransferLayoutFunctor { paddle::platform::MKLDNNDeviceContext::tls() .set_cur_paddle_data_layout(in_layout); } - out_tensor.set_layout(DataLayout::kMKLDNN); - out_tensor.set_format(out_format); + + auto out_tz = phi::vectorize(out_tensor.dims()); + dnnl::memory::data_type in_type = framework::ToMKLDNNDataType( + framework::TransToProtoVarType(in_tensor.dtype())); + + dnnl::memory::desc out_mem_desc(out_tz, in_type, out_format); + out_tensor.set_mem_desc(out_mem_desc); } else { auto target_layout = paddle::platform::MKLDNNDeviceContext::tls() .get_cur_paddle_data_layout(); diff --git a/paddle/phi/kernels/funcs/data_layout_transform.cc b/paddle/phi/kernels/funcs/data_layout_transform.cc index 9d2d0bf3b5..767566cef2 100644 --- a/paddle/phi/kernels/funcs/data_layout_transform.cc +++ b/paddle/phi/kernels/funcs/data_layout_transform.cc @@ -115,8 +115,6 @@ void innerTransDataLayoutFromOneDNN(DataLayout in_layout, out->set_layout(DataLayout::kNCHW); VLOG(10) << "out->layout: " << out->layout() << " in->dims: " << in.dims() << " out->dims: " << out->dims(); - // reset format since the out tensor will be feed to non-MKLDNN OPkernel - out->set_format(OneDNNMemoryFormat::undef); } #endif -- GitLab