247 lines
9.8 KiB
C++
247 lines
9.8 KiB
C++
// 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/fluid/imperative/layout_autotune.h"
|
|
|
|
#include "paddle/common/errors.h"
|
|
#include "paddle/fluid/eager/api/utils/global_utils.h"
|
|
#include "paddle/fluid/framework/op_info.h"
|
|
#include "paddle/fluid/imperative/layout_transformer.h"
|
|
#include "paddle/phi/backends/gpu/gpu_info.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
namespace paddle::imperative {
|
|
|
|
LayoutAutoTune::LayoutAutoTune() {
|
|
const auto& op_info = paddle::framework::OpInfoMap::Instance().map();
|
|
for (const auto& info : op_info) {
|
|
// only when op was not in Lightly、Heavily or Agnostic Set
|
|
if (IsLightlyLayoutSensitive(info.first) ||
|
|
IsHeavilyLayoutSensitive(info.first) || IsLayoutAgnostic(info.first)) {
|
|
VLOG(7) << "Already exists in Layout OP: " << info.first;
|
|
continue;
|
|
}
|
|
|
|
// only record forward operators
|
|
if (info.first.find("_grad") != std::string::npos) {
|
|
continue;
|
|
}
|
|
|
|
auto* attr_checker = info.second.Checker();
|
|
bool layout_agnostic = true;
|
|
if (attr_checker) {
|
|
auto attrs = attr_checker->GetDefaultAttrMap();
|
|
// Attribute name is fuzzy matched, such as start and start_axis.
|
|
for (auto& attr : attrs) {
|
|
auto attr_name = attr.first;
|
|
VLOG(8) << "OP: " << info.first << " Attr Name: " << attr_name;
|
|
if (attr_name.find("axis") != std::string::npos ||
|
|
attr_name.find("axes") != std::string::npos ||
|
|
attr_name.find("dim") != std::string::npos ||
|
|
attr_name.find("start") != std::string::npos ||
|
|
attr_name.find("end") != std::string::npos) {
|
|
VLOG(8) << "Lightly layout sensitive OP: " << info.first;
|
|
layout_agnostic = false;
|
|
lightly_layout_sensitive_ops_.emplace(info.first);
|
|
break;
|
|
}
|
|
}
|
|
|
|
if ((attrs.find("data_format") != attrs.end() ||
|
|
attrs.find("data_layout") != attrs.end()) &&
|
|
layout_agnostic == true) {
|
|
VLOG(8) << "Heavily layout sensitive OP: " << info.first;
|
|
heavily_layout_sensitive_ops_.emplace(info.first);
|
|
layout_agnostic = false; // NOLINT
|
|
continue;
|
|
}
|
|
}
|
|
|
|
// some normalization operators such as instance_norm and layer_norm
|
|
// do not have data_format attr, but are layout sensitive.
|
|
if (info.first.find("norm") != std::string::npos && layout_agnostic) {
|
|
lightly_layout_sensitive_ops_.emplace(info.first);
|
|
continue;
|
|
}
|
|
|
|
if (layout_agnostic) {
|
|
VLOG(8) << "Layout agnostic_ops: " << info.first;
|
|
layout_agnostic_ops_.emplace(info.first);
|
|
}
|
|
}
|
|
|
|
VLOG(6) << "The number of layout agnostic OPs: "
|
|
<< layout_agnostic_ops_.size() << ", heavily layout sensitive OPs: "
|
|
<< heavily_layout_sensitive_ops_.size()
|
|
<< ", lightly layout sensitive OPs: "
|
|
<< lightly_layout_sensitive_ops_.size();
|
|
}
|
|
|
|
template <typename VarType>
|
|
paddle::imperative::NameVarMap<VarType> DealHeavilyLayoutSensitive(
|
|
const std::string& op_type,
|
|
const paddle::imperative::NameVarMap<VarType>& ins,
|
|
const paddle::imperative::NameVarMap<VarType>& outs,
|
|
paddle::framework::AttributeMap* attrs,
|
|
const std::shared_ptr<imperative::Tracer>& tracer) {
|
|
std::shared_ptr<LayoutTransformer<VarType>> transposer = nullptr;
|
|
transposer =
|
|
std::make_shared<HeavilyLayoutSensitiveOpTransformer<VarType>>(op_type);
|
|
transposer->SetArguments(
|
|
{"Input", "X"}, {"Output", "Out", "Y"}, {"data_format", "data_layout"});
|
|
|
|
return transposer->Apply(ins, outs, attrs, tracer);
|
|
}
|
|
|
|
template <typename VarType>
|
|
paddle::imperative::NameVarMap<VarType> DealLightlyLayoutSensitive(
|
|
const std::string& op_type,
|
|
const paddle::imperative::NameVarMap<VarType>& ins,
|
|
const paddle::imperative::NameVarMap<VarType>& outs,
|
|
paddle::framework::AttributeMap* attrs,
|
|
const std::shared_ptr<imperative::Tracer>& tracer) {
|
|
std::shared_ptr<LayoutTransformer<VarType>> transposer = nullptr;
|
|
if (op_type == "transpose2") {
|
|
transposer = std::make_shared<TransposeOpTransformer<VarType>>(op_type);
|
|
} else if (op_type == "flatten_contiguous_range") {
|
|
transposer = std::make_shared<FlattenOpTransformer<VarType>>(op_type);
|
|
} else if (op_type == "arg_max") {
|
|
transposer = std::make_shared<ArgmaxOpTransformer<VarType>>(op_type);
|
|
} else if (op_type == "concat") {
|
|
transposer = std::make_shared<ConcatOpTransformer<VarType>>(op_type);
|
|
} else if (op_type.find("elementwise_") != std::string::npos) {
|
|
transposer = std::make_shared<ElementwiseOpTransformer<VarType>>(op_type);
|
|
} else {
|
|
VLOG(4) << op_type
|
|
<< "'s LayoutTransformer is unimplemented. Use default "
|
|
"LightlyLayoutTransformer instead.";
|
|
transposer =
|
|
std::make_shared<LightlyLayoutSensitiveOpTransformer<VarType>>(op_type);
|
|
}
|
|
return transposer->Apply(ins, outs, attrs, tracer);
|
|
}
|
|
|
|
LayoutAutotuneGuard::LayoutAutotuneGuard(std::shared_ptr<Tracer> tracer,
|
|
bool use_autotune)
|
|
: tracer_(tracer), pre_layout_autotune_(tracer_->UseLayoutAutoTune()) {
|
|
if (pre_layout_autotune_ != use_autotune) {
|
|
tracer_->EnableLayoutAutoTune();
|
|
if (!use_autotune) {
|
|
tracer_->DisableLayoutAutoTune();
|
|
}
|
|
}
|
|
}
|
|
|
|
LayoutAutotuneGuard::~LayoutAutotuneGuard() {
|
|
if (pre_layout_autotune_) { // NOLINT
|
|
tracer_->EnableLayoutAutoTune();
|
|
} else {
|
|
tracer_->DisableLayoutAutoTune();
|
|
}
|
|
}
|
|
|
|
template <typename VarType>
|
|
paddle::imperative::NameVarMap<VarType> AutoTuneLayout(
|
|
const std::string& op_type,
|
|
const paddle::imperative::NameVarMap<VarType>& ins,
|
|
const paddle::imperative::NameVarMap<VarType>& outs,
|
|
paddle::framework::AttributeMap* attrs,
|
|
const std::shared_ptr<imperative::Tracer>& tracer) {
|
|
if (!tracer->UseLayoutAutoTune() ||
|
|
op_type.find("_grad") != std::string::npos) {
|
|
return ins;
|
|
}
|
|
// When layout autotuning is enabled, the tuner will check the desired layout.
|
|
// (1) If the desired layout is undefined, and there is no convolutional
|
|
// layers, layout optimization is unnecessary. Otherwise, the desired layout
|
|
// will be set to the best layout only when these is a convolutional layer
|
|
// with
|
|
// NCHW-Layout and the TensorCore is available.
|
|
// (2) If the desired layout is defined, run the transposer.
|
|
|
|
if (LayoutAutoTune::Instance().GetDesiredLayout() == DataLayout::UNDEFINED) {
|
|
// Layout autotune only supports model with convolutional layers
|
|
if (op_type != "conv2d") {
|
|
return ins;
|
|
} else {
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
if (!phi::backends::gpu::TensorCoreAvailable()) {
|
|
tracer->DisableLayoutAutoTune();
|
|
return ins;
|
|
}
|
|
#endif
|
|
auto conv_in_type = framework::proto::VarType::FP32;
|
|
auto& in_vars = ins.at("Input")[0];
|
|
if (GetDataType<VarType>(in_vars) == framework::proto::VarType::FP16) {
|
|
conv_in_type = framework::proto::VarType::FP16;
|
|
}
|
|
bool is_tune_fp32 =
|
|
(PADDLE_GET_CONST(std::string, (*attrs)["data_format"]) == "NHWC") &&
|
|
(conv_in_type == framework::proto::VarType::FP32);
|
|
bool is_tune_fp16 =
|
|
(PADDLE_GET_CONST(std::string, (*attrs)["data_format"]) == "NCHW") &&
|
|
(conv_in_type == framework::proto::VarType::FP16 ||
|
|
conv_in_type == framework::proto::VarType::BF16);
|
|
if (is_tune_fp32) {
|
|
LayoutAutoTune::Instance().SetDesiredLayout(DataLayout::NCHW);
|
|
LayoutAutoTune::Instance().SetDefaultLayout(DataLayout::NHWC);
|
|
} else if (is_tune_fp16) {
|
|
LayoutAutoTune::Instance().SetDesiredLayout(DataLayout::NHWC);
|
|
LayoutAutoTune::Instance().SetDefaultLayout(DataLayout::NCHW);
|
|
} else {
|
|
tracer->DisableLayoutAutoTune();
|
|
return ins;
|
|
}
|
|
VLOG(3) << "Tune the layout from "
|
|
<< PADDLE_GET_CONST(std::string, (*attrs)["data_format"])
|
|
<< " to "
|
|
<< common::DataLayoutToString(
|
|
LayoutAutoTune::Instance().GetDesiredLayout());
|
|
}
|
|
}
|
|
|
|
if (LayoutAutoTune::Instance().IsHeavilyLayoutSensitive(op_type)) {
|
|
return DealHeavilyLayoutSensitive<VarType>(
|
|
op_type, ins, outs, attrs, tracer);
|
|
} else if (LayoutAutoTune::Instance().IsLightlyLayoutSensitive(op_type)) {
|
|
return DealLightlyLayoutSensitive<VarType>(
|
|
op_type, ins, outs, attrs, tracer);
|
|
} else {
|
|
std::shared_ptr<LayoutTransformer<VarType>> transposer = nullptr;
|
|
if (LayoutAutoTune::Instance().IsLayoutAgnostic(op_type)) {
|
|
transposer = std::make_shared<LayoutTransformer<VarType>>(op_type);
|
|
}
|
|
PADDLE_ENFORCE_NOT_NULL(
|
|
transposer,
|
|
common::errors::Unimplemented(
|
|
"%s 's LayoutTransformer is unimplemented.", op_type));
|
|
return transposer->Apply(ins, outs, attrs, tracer);
|
|
}
|
|
}
|
|
|
|
template paddle::imperative::NameVarMap<VarBase> AutoTuneLayout<VarBase>(
|
|
const std::string& op_type,
|
|
const paddle::imperative::NameVarMap<VarBase>& ins,
|
|
const paddle::imperative::NameVarMap<VarBase>& outs,
|
|
paddle::framework::AttributeMap* attrs,
|
|
const std::shared_ptr<imperative::Tracer>& tracer);
|
|
template paddle::imperative::NameVarMap<egr::EagerVariable>
|
|
AutoTuneLayout<egr::EagerVariable>(
|
|
const std::string& op_type,
|
|
const paddle::imperative::NameVarMap<egr::EagerVariable>& ins,
|
|
const paddle::imperative::NameVarMap<egr::EagerVariable>& outs,
|
|
paddle::framework::AttributeMap* attrs,
|
|
const std::shared_ptr<imperative::Tracer>& tracer);
|
|
|
|
} // namespace paddle::imperative
|