// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. #include "paddle/phi/kernels/conv_transpose_kernel.h" #include "paddle/phi/backends/onednn/onednn_helper.h" #include "paddle/phi/backends/onednn/onednn_reuse.h" #include "paddle/phi/core/compat/get_kerneltype_forvar_utils.h" #include "paddle/phi/core/expect.h" #include "paddle/phi/core/kernel_registry.h" #include "paddle/phi/kernels/cpu/conv_util.h" #include "paddle/phi/kernels/funcs/data_layout_transform.h" namespace phi { struct DeconvolutionCache { dnnl::deconvolution_forward deconvolution_forward; dnnl::memory src_mem; dnnl::memory weights_mem; dnnl::memory bias_mem; dnnl::memory dst_mem; }; inline dnnl::memory::dims GetWeightsTz(const DenseTensor* filter, const int groups) { auto weights_tz = vectorize(filter->dims()); int g = std::max(groups, 1); int g_dim = (g > 1) ? 1 : 0; funcs::GetGroupConvWeightsTz(weights_tz, g); // gIOHW -> gOIHW || IOHW -> OIHW std::swap(weights_tz[g_dim + 0], weights_tz[g_dim + 1]); return weights_tz; } template class ConvTransposeOneDNNHandlerT : public funcs::OneDNNHandlerNoCachingT { private: const bool is_test_; public: ConvTransposeOneDNNHandlerT(const OneDNNContext& dev_ctx, const DenseTensor* x, const DenseTensor* filter, const DenseTensor* bias, const std::vector& strides_in, const std::vector& paddings_in, const std::string& padding_algorithm, int groups, const std::vector& dilations_in, DenseTensor* out) : funcs::OneDNNHandlerNoCachingT( dev_ctx.GetEngine(), dev_ctx.GetPlace()), is_test_(dev_ctx.HasDnnAttr("is_test") ? PADDLE_GET_CONST(bool, dev_ctx.GetDnnAttr("is_test")) : false) { PADDLE_ENFORCE_EQ(is_test_, true, common::errors::InvalidArgument( "ConvTransposeOneDNN works only for inference. " "The attribute \'is_test\' value should be set to " "True, but got is_test=False.")); PADDLE_ENFORCE_EQ( x->layout(), DataLayout::ONEDNN, common::errors::InvalidArgument( "Got wrong layout = %d for Input tensor.", x->layout())); PADDLE_ENFORCE_EQ( filter->layout(), DataLayout::ONEDNN, common::errors::InvalidArgument( "The filter tensor's layout should be %d, but got %d.", DataLayout::ONEDNN, filter->layout())); PADDLE_ENFORCE_EQ( x->dims().size(), 4, common::errors::InvalidArgument("Input must be with 4 dimensions, " "i.e. NCHW. but got dimension =%d", x->dims().size())); PADDLE_ENFORCE_EQ( filter->dims().size(), 4, common::errors::InvalidArgument("Filter must be with 4 dimensions, " "i.e. OIHW, but got dimension =%d", filter->dims().size())); if (bias) { PADDLE_ENFORCE_EQ( bias->layout(), DataLayout::ONEDNN, common::errors::InvalidArgument( "The bias tensor's laytout should be %d, but got %d.", DataLayout::ONEDNN, bias->layout())); PADDLE_ENFORCE_EQ( bias->dims().size(), 1, common::errors::InvalidArgument("Bias must only have 1 dimension, " "i.e. X, but got dimension = %d .", bias->dims().size())); } dnnl::memory::dims strides(begin(strides_in), end(strides_in)); dnnl::memory::dims paddings(begin(paddings_in), end(paddings_in)); dnnl::memory::dims dilations(begin(dilations_in), end(dilations_in)); PADDLE_ENFORCE_EQ( strides.size(), 2, common::errors::Unimplemented( "Now we only support 2d oneDNN convolution transpose op")); const auto x_dims = x->dims(); const auto x_data_dims = slice_ddim(x_dims, 2, x_dims.size()); const auto filter_dims = filter->dims(); const auto filter_data_dims = slice_ddim(filter_dims, 2, filter_dims.size()); const auto ksize = vectorize(filter_data_dims); UpdatePaddingAndDilation( &paddings, &dilations, padding_algorithm, x_data_dims, strides, ksize); std::transform( dilations.begin(), dilations.end(), dilations.begin(), [](int64_t i) { return i - 1; }); const auto src_tz = vectorize(x->dims()); const auto weights_tz = GetWeightsTz(filter, groups); const auto dst_tz = vectorize(out->dims()); const auto onednn_paddings = funcs::ToOneDNNPadding(paddings); /* create memory descriptor for convolution without specified format * ('any') which lets a primitive (convolution in this case) choose * the memory format preferred for best performance */ auto chosen_memory_format = funcs::OneDNNMemoryFormat::any; auto data_type = dnnl::memory::data_type::f32; const bool is_bfloat16 = dev_ctx.HasDnnAttr("mkldnn_data_type") ? PADDLE_GET_CONST(std::string, dev_ctx.GetDnnAttr("mkldnn_data_type")) == "bfloat16" : false; const bool is_onednn_BFLOAT16 = dev_ctx.HasDnnAttr("onednn_data_type") ? PADDLE_GET_CONST(std::string, dev_ctx.GetDnnAttr("onednn_data_type")) == "bfloat16" : is_bfloat16; if (is_onednn_BFLOAT16 || std::is_same::value) { data_type = dnnl::memory::data_type::bf16; } const auto src_md = funcs::OneDNNMemDesc(src_tz, data_type, chosen_memory_format); const auto weights_md = funcs::OneDNNMemDesc(weights_tz, data_type, chosen_memory_format); const auto dst_md = funcs::OneDNNMemDesc( dst_tz, funcs::OneDNNGetDataType(), chosen_memory_format); auto fwd_prop_kind = is_test_ ? dnnl::prop_kind::forward_inference : dnnl::prop_kind::forward_training; if (bias) { std::vector bias_tz = vectorize(bias->dims()); const auto bias_md = funcs::OneDNNMemDesc( bias_tz, data_type, funcs::OneDNNMemoryFormat::x); this->AcquireForwardPrimitiveDescriptor( fwd_prop_kind, dnnl::algorithm::deconvolution_direct, src_md, weights_md, bias_md, dst_md, strides, dilations, onednn_paddings[0], onednn_paddings[1]); } else { this->AcquireForwardPrimitiveDescriptor( fwd_prop_kind, dnnl::algorithm::deconvolution_direct, src_md, weights_md, dst_md, strides, dilations, onednn_paddings[0], onednn_paddings[1]); } } std::shared_ptr AcquireSrcMemoryWithReorder( const DenseTensor* x) { const T* input_data = x->data(); return funcs::OneDNNHandlerNoCachingT:: AcquireMemoryWithReorder(x->mem_desc(), this->fwd_pd_->src_desc(), funcs::to_void_cast(input_data)); } std::shared_ptr AcquireWeightsMemoryWithReorder( const OneDNNContext& dev_ctx, const std::string& key, const DenseTensor* filter, const int& groups) { const K* filter_data = filter->data(); auto weights_tz = GetWeightsTz(filter, groups); int g = std::max(groups, 1); auto user_src_md = funcs::OneDNNMemDesc(weights_tz, funcs::OneDNNGetDataType(), (g == 1) ? funcs::OneDNNMemoryFormat::iohw : funcs::OneDNNMemoryFormat::giohw); return this->template AcquireMemoryWithReorder( dev_ctx, user_src_md, this->fwd_pd_->weights_desc(), funcs::to_void_cast(filter_data), key, "@weights_mem_p", is_test_); } template std::shared_ptr AcquireMemoryWithReorder( const OneDNNContext& dev_ctx, const dnnl::memory::desc& user_md, const dnnl::memory::desc& target_md, void* ptr, const std::string& key, const std::string& suffix, bool is_persistent = false, const std::vector& scale_data = {1.0f}, int mask = 0) { const auto target_key = key + suffix + "_target"; const auto key_reorder_p = key + suffix + "reorder_p"; const auto user_key = key + suffix + "_user"; auto target_memory_p = std::static_pointer_cast(dev_ctx.GetBlob(target_key)); if (target_memory_p == nullptr) { auto user_memory_p = std::make_shared(user_md, this->engine_, ptr); if (user_md != target_md) { target_memory_p = std::make_shared(target_md, this->engine_); dnnl::reorder::primitive_desc reorder_pdesc; if (funcs::is_int8()) { dnnl::primitive_attr attr; attr.set_scales_mask(DNNL_ARG_DST, mask); reorder_pdesc = dnnl::reorder::primitive_desc( *user_memory_p, *target_memory_p, attr); } else { reorder_pdesc = dnnl::reorder::primitive_desc(*user_memory_p, *target_memory_p); } auto reorder_p = std::make_shared(reorder_pdesc); dev_ctx.SetBlob(key_reorder_p, reorder_p); auto& astream = OneDNNContext::tls().get_stream(); std::unordered_map reorder_args; reorder_args.insert({DNNL_ARG_SRC, *user_memory_p}); reorder_args.insert({DNNL_ARG_DST, *target_memory_p}); if (funcs::is_int8()) { auto scale_md = dnnl::memory::desc({static_cast(scale_data.size())}, dnnl::memory::data_type::f32, dnnl::memory::format_tag::x); auto scale_data_mem = dnnl::memory(scale_md, this->engine_); scale_data_mem.set_data_handle( funcs::to_void_cast(scale_data.data())); reorder_args.insert( {DNNL_ARG_ATTR_SCALES | DNNL_ARG_DST, scale_data_mem}); } reorder_p->execute(astream, reorder_args); astream.wait(); } else { target_memory_p = user_memory_p; } dev_ctx.SetBlob(user_key, user_memory_p); dev_ctx.SetBlob(target_key, target_memory_p); } else if (!is_persistent) { auto& astream = OneDNNContext::tls().get_stream(); auto user_memory_p = std::static_pointer_cast(dev_ctx.GetBlob(user_key)); user_memory_p->set_data_handle(ptr); // TODO(jczaja): Here we detect if reorder is cached it means it is needed // need to change this to get rid of keys auto reorder_p = std::static_pointer_cast( dev_ctx.GetBlob(key_reorder_p)); if (reorder_p != nullptr) { reorder_p->execute( astream, {{DNNL_ARG_FROM, *user_memory_p}, {DNNL_ARG_TO, *target_memory_p}}); astream.wait(); } } return target_memory_p; } std::shared_ptr AcquireBiasMemoryWithReorder( const OneDNNContext& dev_ctx, const std::string& key, const DenseTensor* bias) { const K* bias_data = bias->data(); auto user_bias_md = funcs::OneDNNMemDesc(vectorize(bias->dims()), funcs::OneDNNGetDataType(), funcs::OneDNNMemoryFormat::x); return this->AcquireMemoryWithReorder(dev_ctx, user_bias_md, this->fwd_pd_->bias_desc(), funcs::to_void_cast(bias_data), key, "@bias_mem_p", is_test_); } }; template void PrepareSrcMem(const std::shared_ptr& fc_p UNUSED, const std::shared_ptr& src_mem, const DenseTensor* x, const dnnl::engine& engine) { auto x_md = x->mem_desc().reshape(src_mem->get_desc().get_dims()); if (x_md != src_mem->get_desc()) { dnnl::memory x_mem(x_md, engine, funcs::to_void_cast(x->data())); auto reorder_p = dnnl::reorder(x_mem, *src_mem); auto& astream = OneDNNContext::tls().get_stream(); reorder_p.execute(astream, x_mem, *src_mem); astream.wait(); } else { src_mem->set_data_handle(funcs::to_void_cast(x->data())); } } template void Execute(const OneDNNContext& dev_ctx, const DenseTensor* x, const DenseTensor* filter, const DenseTensor* bias, const std::vector& strides, const std::vector& paddings, const std::string& padding_algorithm, int groups, const std::vector& dilations, DenseTensor* out) { std::shared_ptr conv_p; std::shared_ptr src_memory_p; std::shared_ptr weights_memory_p; std::shared_ptr bias_memory_p; std::shared_ptr dst_memory_p; std::unordered_map args; // Note(ZKK): // Add thread_id to cache_key // fix issue https://github.com/PaddlePaddle/PaddleOCR/issues/15621 // https://github.com/PaddlePaddle/PaddleOCR/issues/15393 std::string cache_key = funcs::CreateKey(dev_ctx, funcs::ThreadIDasStr(), dev_ctx.GetInputsName("Input")[0], dev_ctx.GetInputsName("Filter")[0], vectorize(x->dims()), vectorize(filter->dims())); const auto& onednn_engine = dev_ctx.GetEngine(); auto deconvolution_cache = std::static_pointer_cast(dev_ctx.GetBlob(cache_key)); if (deconvolution_cache) { conv_p = std::make_shared( deconvolution_cache->deconvolution_forward); src_memory_p = std::make_shared(deconvolution_cache->src_mem); PrepareSrcMem(conv_p, src_memory_p, x, onednn_engine); weights_memory_p = std::make_shared(deconvolution_cache->weights_mem); dst_memory_p = std::make_shared(deconvolution_cache->dst_mem); auto out_ptr = dev_ctx.template Alloc(out, dst_memory_p->get_desc().get_size()); dst_memory_p->set_data_handle(out_ptr); args.insert({DNNL_ARG_SRC, *src_memory_p}); args.insert({DNNL_ARG_WEIGHTS, *weights_memory_p}); args.insert({DNNL_ARG_DST, *dst_memory_p}); if (bias) { bias_memory_p = std::make_shared(deconvolution_cache->bias_mem); args.insert({DNNL_ARG_BIAS, *bias_memory_p}); } } else { // Check if bias obey the rules if (bias) { PADDLE_ENFORCE_EQ( bias->layout(), DataLayout::ONEDNN, common::errors::InvalidArgument( "The Bias tensor's layout should be %d, but got %d.", DataLayout::ONEDNN, bias->layout())); PADDLE_ENFORCE_EQ( bias->dims().size(), 1, common::errors::InvalidArgument("Bias must only have 1 dimension, " "i.e. X, but got dimension = %d .", bias->dims().size())); } // Caching Key for weights is needed std::string key = funcs::CreateKey(dev_ctx, dev_ctx.GetInputsName("Input")[0], dev_ctx.GetInputsName("Filter")[0], (bias ? dev_ctx.GetInputsName("Bias")[0] : "")); ConvTransposeOneDNNHandlerT handler(dev_ctx, x, filter, bias, strides, paddings, padding_algorithm, groups, dilations, out); src_memory_p = handler.AcquireSrcMemoryWithReorder(x); key = funcs::ExtendKeyWithThreadInfoIfNeeded(dev_ctx, key); weights_memory_p = handler.AcquireWeightsMemoryWithReorder(dev_ctx, key, filter, groups); dst_memory_p = handler.template AcquireDstMemory(out); conv_p = handler.AcquireForwardPrimitive(); args.insert({DNNL_ARG_SRC, *src_memory_p}); args.insert({DNNL_ARG_WEIGHTS, *weights_memory_p}); args.insert({DNNL_ARG_DST, *dst_memory_p}); if (bias) { bias_memory_p = handler.AcquireBiasMemoryWithReorder(dev_ctx, key, bias); args.insert({DNNL_ARG_BIAS, *bias_memory_p}); } auto cache = std::make_shared(); cache->deconvolution_forward = *conv_p; cache->src_mem = *src_memory_p; cache->weights_mem = *weights_memory_p; cache->dst_mem = *dst_memory_p; if (bias) { cache->bias_mem = *bias_memory_p; } dev_ctx.SetBlob(cache_key, cache); } auto& astream = OneDNNContext::tls().get_stream(); conv_p->execute(astream, args); astream.wait(); out->set_mem_desc(dst_memory_p->get_desc()); } template void Conv2dTransposeKernel(const Context& dev_ctx, const DenseTensor& x, const DenseTensor& filter, const std::vector& strides, const std::vector& paddings, const std::vector& output_padding UNUSED, const IntArray& output_size UNUSED, const std::string& padding_algorithm, int groups, const std::vector& dilations, const std::string& data_format UNUSED, DenseTensor* out) { const bool is_bfloat16 = dev_ctx.HasDnnAttr("mkldnn_data_type") ? PADDLE_GET_CONST(std::string, dev_ctx.GetDnnAttr("mkldnn_data_type")) == "bfloat16" : false; const bool is_onednn_BFLOAT16 = dev_ctx.HasDnnAttr("onednn_data_type") ? PADDLE_GET_CONST(std::string, dev_ctx.GetDnnAttr("onednn_data_type")) == "bfloat16" : is_bfloat16; const bool force_fp32_output = dev_ctx.HasDnnAttr("force_fp32_output") ? PADDLE_GET_CONST(bool, dev_ctx.GetDnnAttr("force_fp32_output")) : false; const bool use_bfloat16 = (!force_fp32_output && is_onednn_BFLOAT16); if (use_bfloat16) { Execute(dev_ctx, &x, &filter, nullptr, strides, paddings, padding_algorithm, groups, dilations, out); } else { Execute(dev_ctx, &x, &filter, nullptr, strides, paddings, padding_algorithm, groups, dilations, out); } } template void Conv2dTransposeBiasKernel(const Context& dev_ctx, const DenseTensor& x, const DenseTensor& filter, const optional& bias, const std::vector& strides, const std::vector& paddings, const std::vector& output_padding UNUSED, const IntArray& output_size UNUSED, const std::string& padding_algorithm, int groups, const std::vector& dilations, const std::string& data_format UNUSED, DenseTensor* out) { const bool is_bfloat16 = dev_ctx.HasDnnAttr("mkldnn_data_type") ? PADDLE_GET_CONST(std::string, dev_ctx.GetDnnAttr("mkldnn_data_type")) == "bfloat16" : false; const bool is_one_BFLOAT16 = dev_ctx.HasDnnAttr("onednn_data_type") ? PADDLE_GET_CONST(std::string, dev_ctx.GetDnnAttr("onednn_data_type")) == "bfloat16" : is_bfloat16; const bool force_fp32_output = dev_ctx.HasDnnAttr("force_fp32_output") ? PADDLE_GET_CONST(bool, dev_ctx.GetDnnAttr("force_fp32_output")) : false; const bool use_bfloat16 = (!force_fp32_output && is_one_BFLOAT16); if (use_bfloat16) { Execute(dev_ctx, &x, &filter, bias.get_ptr(), strides, paddings, padding_algorithm, groups, dilations, out); } else { Execute(dev_ctx, &x, &filter, bias.get_ptr(), strides, paddings, padding_algorithm, groups, dilations, out); } } KernelKey ConvTransposeGetKernelTypeForVar( const GetKernelTypeForVarContext* ctx) { const std::string& var_name = ctx->GetVarName(); const DenseTensor& tensor = ctx->GetTensor(); const KernelKey& expected_kernel_type = ctx->GetKernelKey(); const AttributeMap& attrs = ctx->GetAttrs(); // Only input require reshaping, weights and // bias are having shape in NCHW order if ((var_name == "Input") && (expected_kernel_type.layout() == DataLayout::ONEDNN) && (tensor.layout() != DataLayout::ONEDNN)) { auto it = attrs.find("data_format"); const std::string data_format = PADDLE_GET_CONST(std::string, it->second); auto dl = StringToDataLayout(data_format); // Some models may have intentionally set "AnyLayout" for pool // op. Treat this as NCHW (default data_format value) if (dl != DataLayout::ANY) { return phi::KernelKey(tensor.place(), dl, expected_kernel_type.dtype()); } } return phi::KernelKey( tensor.place(), tensor.layout(), expected_kernel_type.dtype()); } } // namespace phi PD_REGISTER_KERNEL(conv2d_transpose, OneDNN, ONEDNN, phi::Conv2dTransposeKernel, float, phi::bfloat16) { kernel->get_kerneltype_forvar_fn_ = phi::ConvTransposeGetKernelTypeForVar; } PD_REGISTER_KERNEL(conv2d_transpose_bias, OneDNN, ONEDNN, phi::Conv2dTransposeBiasKernel, float, phi::bfloat16) { kernel->get_kerneltype_forvar_fn_ = phi::ConvTransposeGetKernelTypeForVar; }