From 6f5e78265ecc9213b578dfdb155ab6cd19986bd6 Mon Sep 17 00:00:00 2001 From: HongyuJia Date: Tue, 25 Oct 2022 15:22:03 +0800 Subject: [PATCH] [Kernel Selection] Remove hard code of PADDLE_WITH_MKLDNN (Part2 add dnn_fallback flag) (#47200) * use dnn_fallback flag to delete mkldnn hardcode * polish code style * fix protected error * fix const error * fix reduce_op fallback * fix pool_op fallback * add Set function of dnn_fallback_ --- paddle/fluid/framework/operator.cc | 30 +++++++------- paddle/fluid/framework/operator.h | 8 ++++ paddle/fluid/imperative/prepared_operator.cc | 9 +++-- paddle/fluid/operators/cast_op.cc | 28 ++++--------- paddle/fluid/operators/layer_norm_op.cc | 14 +++---- paddle/fluid/operators/optimizers/sgd_op.cc | 31 +++++++-------- paddle/fluid/operators/pool_op.cc | 26 ++++++------- paddle/fluid/operators/reduce_ops/reduce_op.h | 39 +++++++------------ paddle/fluid/operators/sum_op.cc | 29 +++++++------- paddle/fluid/platform/mkldnn_op_list.h | 18 ++------- 10 files changed, 98 insertions(+), 134 deletions(-) diff --git a/paddle/fluid/framework/operator.cc b/paddle/fluid/framework/operator.cc index 14773c0433..5f03a504d4 100644 --- a/paddle/fluid/framework/operator.cc +++ b/paddle/fluid/framework/operator.cc @@ -1388,12 +1388,13 @@ bool OperatorWithKernel::SupportsKernelType( #endif // NOTE(jiahongyu): If MKLDNN can be used, the function SupportsKernelType needs -// to check whether current op supports MKLDNN kernel. There are two statements -// in if condition: -// 1. Whether this op has specific implementation; -// 2. Whether mkldnn kernel can be used. +// to check whether current op supports MKLDNN kernel. There are three +// statements in if condition: +// 1. Whether mkldnn kernel fallbacks to plain kernel; +// 2. Whether this op has specific implementation; +// 3. Whether mkldnn kernel can be used. #ifdef PADDLE_WITH_MKLDNN - if (!paddle::platform::in_mkldnn_white_list(type_) && + if (!this->DnnFallback() && !paddle::platform::in_mkldnn_white_list(type_) && this->CanMKLDNNBeUsed(exe_ctx, kernel_type.data_type_)) { auto tmp_kernel_type = kernel_type; tmp_kernel_type.library_type_ = framework::LibraryType::kMKLDNN; @@ -1569,11 +1570,13 @@ void OperatorWithKernel::RunImpl(const Scope& scope, // NOTE(jiahongyu): The registered MKLDNN kernel have library_type = // LibraryType::kMKLDNN and data_layout_ = DataLayout::kMKLDNN. But the default // values are kPlain, so we need to modify the library_type and data_layout_ -// here. There are two statements in if condition: -// 1. Whether this op has specific implementation; -// 2. Whether mkldnn kernel can be used. +// here. There are three statements in if condition: +// 1. Whether mkldnn kernel fallbacks to plain kernel; +// 2. Whether this op has specific implementation; +// 3. Whether mkldnn kernel can be used. #ifdef PADDLE_WITH_MKLDNN - if (!paddle::platform::in_mkldnn_white_list(type_) && + if (!this->DnnFallback() && + !paddle::platform::in_mkldnn_white_list(type_) && this->CanMKLDNNBeUsed(exe_ctx, kernel_type_->data_type_)) { kernel_type_->library_type_ = framework::LibraryType::kMKLDNN; kernel_type_->data_layout_ = framework::DataLayout::kMKLDNN; @@ -1810,12 +1813,13 @@ OpKernelType OperatorWithKernel::InnerGetExpectedKernelType( // NOTE(jiahongyu): PADDLE_WITH_MKLDNN codes are moved outside function // GetExpectedKernelType, so that if MKLDNN can be used, the library_type_ and -// data_layout_ of expected_kernel_key need to be adjusted. There are two +// data_layout_ of expected_kernel_key need to be adjusted. There are three // statements in if condition: -// 1. Whether this op has specific implementation; -// 2. Whether mkldnn kernel can be used. +// 1. Whether mkldnn kernel fallbacks to plain kernel; +// 2. Whether this op has specific implementation; +// 3. Whether mkldnn kernel can be used. #ifdef PADDLE_WITH_MKLDNN - if (!paddle::platform::in_mkldnn_white_list(type_) && + if (!this->DnnFallback() && !paddle::platform::in_mkldnn_white_list(type_) && this->CanMKLDNNBeUsed(ctx, expected_kernel_key.data_type_)) { expected_kernel_key.library_type_ = framework::LibraryType::kMKLDNN; expected_kernel_key.data_layout_ = framework::DataLayout::kMKLDNN; diff --git a/paddle/fluid/framework/operator.h b/paddle/fluid/framework/operator.h index af629bf6df..5acec948ec 100644 --- a/paddle/fluid/framework/operator.h +++ b/paddle/fluid/framework/operator.h @@ -704,6 +704,10 @@ class OperatorWithKernel : public OperatorBase { kernel_type_.reset(kernel_type); } + bool DnnFallback() const { return dnn_fallback_; } + + void SetDnnFallback(bool dnn_fallback) const { dnn_fallback_ = dnn_fallback; } + private: void RunImpl(const Scope& scope, const platform::Place& place) const final; void RunImpl(const Scope& scope, @@ -756,6 +760,10 @@ class OperatorWithKernel : public OperatorBase { mutable bool all_kernels_must_compute_runtime_shape_ = false; mutable std::mutex cache_update_mutex_; mutable bool enable_cache_transfer_scope_ = false; + // NOTE(jiahongyu): Whether fallback to plain kernel after calling + // GetExpectedKernelType, use this bool flag to solve mkldnn and cudnn hard + // code + mutable bool dnn_fallback_ = false; // NOTE(chenweihang): Similar op members are used to adapt to // new phi kernel, if there is a better design in the future, // we may polish the implementation here diff --git a/paddle/fluid/imperative/prepared_operator.cc b/paddle/fluid/imperative/prepared_operator.cc index 28276ddbf8..c6cc9befbd 100644 --- a/paddle/fluid/imperative/prepared_operator.cc +++ b/paddle/fluid/imperative/prepared_operator.cc @@ -192,11 +192,12 @@ PreparedOp PrepareImpl( // NOTE(jiahongyu): The registered MKLDNN kernel have library_type = // LibraryType::kMKLDNN and data_layout_ = DataLayout::kMKLDNN. But the default // values are kPlain, so we need to modify the library_type and data_layout_ -// here. There are two statements in if condition: -// 1. Whether this op has specific implementation; -// 2. Whether mkldnn kernel can be used. +// here. There are three statements in if condition: +// 1. Whether mkldnn kernel fallbacks to plain kernel; +// 2. Whether this op has specific implementation; +// 3. Whether mkldnn kernel can be used. #ifdef PADDLE_WITH_MKLDNN - if (!paddle::platform::in_mkldnn_white_list(op.Type()) && + if (!op.DnnFallback() && !paddle::platform::in_mkldnn_white_list(op.Type()) && op.CanMKLDNNBeUsed(dygraph_exe_ctx, expected_kernel_key.data_type_)) { expected_kernel_key.library_type_ = framework::LibraryType::kMKLDNN; expected_kernel_key.data_layout_ = framework::DataLayout::kMKLDNN; diff --git a/paddle/fluid/operators/cast_op.cc b/paddle/fluid/operators/cast_op.cc index b994c80283..5505d3b4e3 100644 --- a/paddle/fluid/operators/cast_op.cc +++ b/paddle/fluid/operators/cast_op.cc @@ -91,31 +91,19 @@ class CastOp : public framework::OperatorWithKernel { ctx.device_context()); } -#ifdef PADDLE_WITH_MKLDNN + // NOTE(jiahongyu): Below codes originally enclosed by PADDLE_WITH_MKLDNN int in_dtype = ctx.Attr("in_dtype"); int out_dtype = ctx.Attr("out_dtype"); - auto MKLDNNSupportsCast = [&]() -> bool { - int dtype_fp32 = static_cast(framework::proto::VarType::FP32); - int dtype_bf16 = static_cast(framework::proto::VarType::BF16); + int dtype_fp32 = static_cast(framework::proto::VarType::FP32); + int dtype_bf16 = static_cast(framework::proto::VarType::BF16); - if ((in_dtype != dtype_fp32 && in_dtype != dtype_bf16) || - (out_dtype != dtype_fp32 && out_dtype != dtype_bf16)) - return false; - - return true; - }; - - if (this->CanMKLDNNBeUsed( - ctx, framework::TransToProtoVarType(tensor->dtype())) && - MKLDNNSupportsCast()) { - return framework::OpKernelType( - framework::TransToProtoVarType(tensor->dtype()), - ctx.GetPlace(), - phi::DataLayout::kMKLDNN, - framework::LibraryType::kMKLDNN); + if ((in_dtype != dtype_fp32 && in_dtype != dtype_bf16) || + (out_dtype != dtype_fp32 && out_dtype != dtype_bf16)) { + this->SetDnnFallback(true); } -#endif + // NOTE(jiahongyu): Above codes originally enclosed by PADDLE_WITH_MKLDNN + #ifdef PADDLE_WITH_MLU auto src_type = static_cast(ctx.Attr("in_dtype")); auto dst_type = static_cast(ctx.Attr("out_dtype")); diff --git a/paddle/fluid/operators/layer_norm_op.cc b/paddle/fluid/operators/layer_norm_op.cc index 9d89d57831..8b7a9492d0 100644 --- a/paddle/fluid/operators/layer_norm_op.cc +++ b/paddle/fluid/operators/layer_norm_op.cc @@ -111,17 +111,13 @@ class LayerNormOp : public framework::OperatorWithKernel { const framework::ExecutionContext &ctx) const override { auto input_data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X"); -#ifdef PADDLE_WITH_MKLDNN + // NOTE(jiahongyu): Below codes originally enclosed by PADDLE_WITH_MKLDNN int begin_norm_axis = ctx.Attr("begin_norm_axis"); - if (this->CanMKLDNNBeUsed(ctx, input_data_type) && - begin_norm_axis == - ctx.Input("X")->dims().size() - 1) { - return framework::OpKernelType(input_data_type, - ctx.GetPlace(), - phi::DataLayout::kMKLDNN, - framework::LibraryType::kMKLDNN); + if (begin_norm_axis != + ctx.Input("X")->dims().size() - 1) { + this->SetDnnFallback(true); } -#endif + // NOTE(jiahongyu): Above codes originally enclosed by PADDLE_WITH_MKLDNN return framework::OpKernelType(input_data_type, ctx.GetPlace()); } diff --git a/paddle/fluid/operators/optimizers/sgd_op.cc b/paddle/fluid/operators/optimizers/sgd_op.cc index 099d35502c..72c90aaf49 100644 --- a/paddle/fluid/operators/optimizers/sgd_op.cc +++ b/paddle/fluid/operators/optimizers/sgd_op.cc @@ -35,25 +35,20 @@ class SGDOp : public framework::OperatorWithKernel { const framework::ExecutionContext &ctx) const override { auto data_type = OperatorWithKernel::IndicateVarDataType(ctx, "Param"); -#ifdef PADDLE_WITH_MKLDNN - using dnnl::memory; - if (this->CanMKLDNNBeUsed(ctx, data_type)) { - const auto *param_var = ctx.InputVar("Param"); - const auto *grad_var = ctx.InputVar("Grad"); - - // supported cases - bool dense_param_sparse_grad = param_var->IsType() && - grad_var->IsType(); - bool dense_param_and_grad = param_var->IsType() && - grad_var->IsType(); - - if (dense_param_sparse_grad || dense_param_and_grad) - return framework::OpKernelType(data_type, - ctx.GetPlace(), - phi::DataLayout::kMKLDNN, - framework::LibraryType::kMKLDNN); + // NOTE(jiahongyu): Below codes originally enclosed by PADDLE_WITH_MKLDNN + const auto *param_var = ctx.InputVar("Param"); + const auto *grad_var = ctx.InputVar("Grad"); + + // supported cases + bool dense_param_sparse_grad = param_var->IsType() && + grad_var->IsType(); + bool dense_param_and_grad = param_var->IsType() && + grad_var->IsType(); + if (!(dense_param_sparse_grad || dense_param_and_grad)) { + this->SetDnnFallback(true); } -#endif + // NOTE(jiahongyu): Above codes originally enclosed by PADDLE_WITH_MKLDNN + return framework::OpKernelType(data_type, ctx.device_context()); } diff --git a/paddle/fluid/operators/pool_op.cc b/paddle/fluid/operators/pool_op.cc index 5c8da74817..7842de9b17 100644 --- a/paddle/fluid/operators/pool_op.cc +++ b/paddle/fluid/operators/pool_op.cc @@ -33,6 +33,9 @@ bool CanMKLDNNSupportPool(const framework::ExecutionContext& ctx) { if (ctx.Attr("adaptive") == false) return true; // (jczaja): oneDNN is supporting only unchangable in size pool window auto src_tz = phi::vectorize(ctx.Input("X")->dims()); + if (!ctx.HasAttr("ksize")) { + return false; + } std::vector ksize = ctx.Attr>("ksize"); // Fast but not exhustive check return ((src_tz[src_tz.size() - 1] % ksize[1] == 0) && @@ -50,13 +53,10 @@ framework::OpKernelType PoolOp::GetExpectedKernelType( library_ = framework::LibraryType::kCUDNN; } #endif -#ifdef PADDLE_WITH_MKLDNN - if (library_ == framework::LibraryType::kPlain && - this->CanMKLDNNBeUsed(ctx, data_type) && CanMKLDNNSupportPool(ctx)) { - library_ = framework::LibraryType::kMKLDNN; - layout_ = phi::DataLayout::kMKLDNN; - } -#endif + + // NOTE(jiahongyu): Below codes originally enclosed by PADDLE_WITH_MKLDNN + this->SetDnnFallback(!CanMKLDNNSupportPool(ctx)); + // NOTE(jiahongyu) END: Above codes originally enclosed by PADDLE_WITH_MKLDNN return framework::OpKernelType(data_type, ctx.GetPlace(), layout_, library_); } @@ -95,14 +95,10 @@ framework::OpKernelType PoolOpGrad::GetExpectedKernelType( library_ = framework::LibraryType::kCUDNN; } #endif -#ifdef PADDLE_WITH_MKLDNN - if (library_ == framework::LibraryType::kPlain && - this->CanMKLDNNBeUsed(ctx, input_data_type) && - CanMKLDNNSupportPool(ctx)) { - library_ = framework::LibraryType::kMKLDNN; - layout_ = phi::DataLayout::kMKLDNN; - } -#endif + + // NOTE(jiahongyu): Below codes originally enclosed by PADDLE_WITH_MKLDNN + this->SetDnnFallback(!CanMKLDNNSupportPool(ctx)); + // NOTE(jiahongyu): Above codes originally enclosed by PADDLE_WITH_MKLDNN return framework::OpKernelType( input_data_type, ctx.GetPlace(), layout_, library_); diff --git a/paddle/fluid/operators/reduce_ops/reduce_op.h b/paddle/fluid/operators/reduce_ops/reduce_op.h index 91db65c0e1..027a787cbf 100644 --- a/paddle/fluid/operators/reduce_ops/reduce_op.h +++ b/paddle/fluid/operators/reduce_ops/reduce_op.h @@ -559,6 +559,10 @@ class ReduceOp : public framework::OperatorWithKernel { experimental::DataType::BFLOAT16) return true; + if (!ctx.HasAttr("dim") || !ctx.HasAttr("reduce_all")) { + return false; + } + auto reduce_dims = ctx.Attr>("dim"); const bool reduce_all = ctx.Attr("reduce_all"); int ndims = ctx.Input("X")->dims().size(); @@ -586,18 +590,12 @@ class ReduceOp : public framework::OperatorWithKernel { // choose cudnn kernel if the runtime supported. auto input_data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X"); - if (ctx.Input("X")->dims().size() > 5) - return framework::OpKernelType(input_data_type, ctx.GetPlace()); - -#ifdef PADDLE_WITH_MKLDNN - if (this->CanMKLDNNBeUsed(ctx, input_data_type) && - HasOptimizedOneDNNKernel(ctx)) { - return framework::OpKernelType(input_data_type, - ctx.GetPlace(), - phi::DataLayout::kMKLDNN, - framework::LibraryType::kMKLDNN); + // NOTE(jiahongyu): Below codes originally enclosed by PADDLE_WITH_MKLDNN + if (ctx.Input("X")->dims().size() > 5 || + !HasOptimizedOneDNNKernel(ctx)) { + this->SetDnnFallback(true); } -#endif + // NOTE(jiahongyu): Above codes originally enclosed by PADDLE_WITH_MKLDNN if (input_data_type == framework::proto::VarType::FP16) { PADDLE_ENFORCE_EQ( @@ -674,22 +672,13 @@ class ReduceGradOp : public framework::OperatorWithKernel { ? static_cast(out_dtype) : OperatorWithKernel::IndicateVarDataType( ctx, framework::GradVarName("Out")); -#ifdef PADDLE_WITH_MKLDNN - auto CanMKLDNNReduceGradBeUsed = [&]() { - auto dx_dims = ctx.Input("X")->dims(); - if (dx_dims.size() > 5) return false; // max 5D tensor is supported - - return true; - }; - if (this->CanMKLDNNBeUsed(ctx, input_data_type) && - CanMKLDNNReduceGradBeUsed()) { - return framework::OpKernelType(input_data_type, - ctx.GetPlace(), - phi::DataLayout::kMKLDNN, - framework::LibraryType::kMKLDNN); + // NOTE(jiahongyu): Below codes originally enclosed by PADDLE_WITH_MKLDNN + // max 5D tensor is supported + if (ctx.Input("X")->dims().size() > 5) { + dnn_fallback_ = true; } -#endif + // NOTE(jiahongyu): Above codes originally enclosed by PADDLE_WITH_MKLDNN return framework::OpKernelType(input_data_type, ctx.GetPlace()); } diff --git a/paddle/fluid/operators/sum_op.cc b/paddle/fluid/operators/sum_op.cc index 242f3c6f7f..5828e681f2 100644 --- a/paddle/fluid/operators/sum_op.cc +++ b/paddle/fluid/operators/sum_op.cc @@ -76,22 +76,21 @@ class SumOp : public framework::OperatorWithKernel { "Sum operator should have at least one tensor")); auto data_type = static_cast(dtype); -#ifdef PADDLE_WITH_MKLDNN - if (this->CanMKLDNNBeUsed(ctx, data_type) && - (data_type == framework::proto::VarType::FP32 || - data_type == framework::proto::VarType::BF16) && - ctx.OutputVar("Out")->IsType()) { - if (std::all_of( - x_vars.begin(), x_vars.end(), [](const framework::Variable* v) { - return v->IsType(); - })) { - return framework::OpKernelType(data_type, - ctx.GetPlace(), - phi::DataLayout::kMKLDNN, - framework::LibraryType::kMKLDNN); - } + + // NOTE(jiahongyu): Below codes originally enclosed by PADDLE_WITH_MKLDNN + if (!((data_type == framework::proto::VarType::FP32 || + data_type == framework::proto::VarType::BF16) && + ctx.OutputVar("Out")->IsType())) { + this->SetDnnFallback(true); + } else if (!std::all_of(x_vars.begin(), + x_vars.end(), + [](const framework::Variable* v) { + return v->IsType(); + })) { + this->SetDnnFallback(true); } -#endif + // NOTE(jiahongyu): Above codes originally enclosed by PADDLE_WITH_MKLDNN + return framework::OpKernelType(data_type, ctx.GetPlace()); } else if (x_vars[0]->IsType()) { for (auto& var : x_vars) { diff --git a/paddle/fluid/platform/mkldnn_op_list.h b/paddle/fluid/platform/mkldnn_op_list.h index 35046bcd9c..686a70b1cf 100644 --- a/paddle/fluid/platform/mkldnn_op_list.h +++ b/paddle/fluid/platform/mkldnn_op_list.h @@ -27,18 +27,14 @@ namespace platform { // TODO(jiahongyu): Delete mkldnn_white_list and fully support // PADDLE_WITH_MKLDNN of GetExpectedKernelType. static const std::unordered_set mkldnn_white_list = { - "cast", - "transfer_dtype", - "layer_norm", + // NOTE(jiahongyu): Below ops use mem_desc function, which is encoded by + // PADDLE_WITH_MKLDNN in DenseTensor. The hardcodes within + // GetExpectedKernelType of these ops cannot be deleted now. "pad2d", "pad3d", - "pool2d", - "pool2d_grad", "slice", "slice_grad", "split", - "sum", - "sgd", // NOTE(jiahongyu): squeeze MKLDNN kernel are disabled // (https://github.com/PaddlePaddle/Paddle/pull/35781). If these MKLDNN // kernels and codes are deleted in the future, attributes `use_mkldnn` @@ -59,14 +55,6 @@ static const std::unordered_set mkldnn_white_list = { "flatten_grad", "flatten2", "flatten2_grad", - // NOTE(jiahongyu): After fixing GetExpectedKernelType in ReduceOp, reduce - // series hard code can be deleted together. - "reduce_max", - "reduce_mean", - "reduce_mean_grad", - "reduce_min", - "reduce_sum", - "reduce_sum_grad", // NOTE(jiahongyu): Below ops register kernel with customized_type_value, we // need to analysis and solve them one-by-one. "prior_box"}; -- GitLab