Files
paddlepaddle--paddle/paddle/fluid/imperative/op_base.h
T
2026-07-13 12:40:42 +08:00

323 lines
9.9 KiB
C++

// Copyright (c) 2020 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.
#pragma once
#include <atomic>
#include <map>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "paddle/fluid/framework/type_defs.h"
#include "paddle/fluid/imperative/saved_variable_wrapper_list.h"
#include "paddle/fluid/imperative/type_defs.h"
#include "paddle/fluid/imperative/variable_wrapper.h"
#include "paddle/phi/common/place.h"
namespace paddle {
namespace imperative {
const static framework::AttributeMap empty_default_attr_map; // NOLINT
// TODO(zjl): to support py_func layer
class OpBase {
public:
OpBase() = default;
OpBase(const OpBase&) = delete;
OpBase(OpBase&&) = default;
OpBase& operator=(const OpBase&) = delete;
OpBase& operator=(OpBase&&) = default;
~OpBase() { VLOG(3) << "Destruct Op: " << Type(); }
const std::string& Type() const {
return op_ ? op_->Type() : UnknownOpType();
}
const framework::AttributeMap& Attrs() const { return attrs_; }
const framework::AttributeMap& DefaultAttrsMap() const {
return *default_attrs_;
}
const framework::OpInfo& Info() const {
PADDLE_ENFORCE_NOT_NULL(op_,
common::errors::PreconditionNotMet(
"OpBase::Info() should be called after "
"OpBase::SetType() is called"));
return op_->Info();
}
const framework::OperatorBase& InnerOp() const {
PADDLE_ENFORCE_NOT_NULL(op_,
common::errors::PreconditionNotMet(
"OpBase::InnerOp() should be called after "
"OpBase::SetType() is called"));
return *op_;
}
PADDLE_API void ClearBackwardTrace();
NameVarMap<VariableWrapper>* GetMutableOutsMap() { return &outs_; }
NameVarMap<VariableWrapper>* GetMutableInsMap() { return &ins_; }
const NameVarMap<VariableWrapper>& GetInsMap() const { return ins_; }
const NameVarMap<VariableWrapper>& GetOutsMap() const { return outs_; }
PADDLE_API void SetType(const std::string& type);
void CheckAttrs() {
auto& info = Info();
if (info.Checker() != nullptr) {
info.Checker()->Check(&attrs_, true);
}
}
void SetInput(const std::string& name,
VariableWrapperList vars,
bool is_grad) {
auto& in_vars = ins_[name];
*(in_vars.MutableVarList()) = std::move(vars);
in_vars.SetIsGrad(is_grad);
}
void SetOutput(const std::string& name,
VariableWrapperList vars,
bool is_grad) {
auto& out_vars = outs_[name];
*(out_vars.MutableVarList()) = std::move(vars);
out_vars.SetIsGrad(is_grad);
}
void SetAttrMap(const framework::AttributeMap& attrs) { attrs_ = attrs; }
void SetDefaultAttrsMap(const framework::AttributeMap& default_attrs) {
default_attrs_ = &default_attrs;
}
void SetAttr(const std::string& name, const framework::Attribute& v) {
attrs_[name] = v;
}
void SetBlockAttr(const std::string& name UNUSED,
framework::BlockDesc* block UNUSED) {
PADDLE_THROW(common::errors::PermissionDenied(
"SetBlockAttr is not support in dygraph OpBase"));
}
const framework::AttributeMap& Attrs() { return attrs_; }
const framework::AttributeMap& DefaultAttrsMap() {
if (default_attrs_ == nullptr) {
return empty_default_attr_map;
}
return *default_attrs_;
}
bool HasAttr(const std::string& name) const {
VLOG(6) << "Default attrs: " << default_attrs_;
VLOG(6) << "attrs: " << &attrs_;
return attrs_.count(name) > 0 || default_attrs_->count(name) > 0;
}
const framework::Attribute& GetAttr(const std::string& name) const {
auto it = attrs_.find(name);
if (it != attrs_.end()) {
return it->second;
} else {
auto it_default = default_attrs_->find(name);
PADDLE_ENFORCE_NE(
it_default,
default_attrs_->end(),
common::errors::NotFound("can not find attribute [%s]", name));
return it_default->second;
}
}
template <typename T>
inline const T& Attr(const std::string& name) const {
return PADDLE_GET_CONST(T, GetAttr(name));
}
size_t id() const { return id_; }
void SetId(size_t id) { id_ = id; }
const phi::Place& place() const { return place_; }
void SetPlace(const phi::Place& place) { place_ = place; }
void EnforceHasInOut() const {
PADDLE_ENFORCE_NE(
ins_.empty() && outs_.empty(),
true,
common::errors::NotFound(
"Inputs and outputs of %s do not exist. This may be because:\n"
"1. You use some output variables of the previous batch as the "
"inputs of the current batch. Please try to call \"stop_gradient "
"= True\" or \"detach()\" for these variables.\n"
"2. You calculate backward twice for the same subgraph without "
"setting retain_graph=True. Please set retain_graph=True in the "
"first backward call.\n\n",
Type()));
}
static size_t GenerateUniqueId() {
static std::atomic<size_t> unique_id{0};
return unique_id.fetch_add(1);
}
static void Run(const framework::OperatorBase& op,
const NameVarMap<VarBase>& ins,
const NameVarMap<VarBase>& outs,
const framework::AttributeMap& attrs,
const framework::AttributeMap& default_attrs,
const phi::Place& place);
static void Run(const framework::OperatorBase& op,
const NameVarMap<VariableWrapper>& ins,
const NameVarMap<VariableWrapper>& outs,
const framework::AttributeMap& attrs,
const framework::AttributeMap& default_attrs,
const phi::Place& place);
static void Run(const framework::OperatorBase& op,
const NameVarMap<egr::EagerVariable>& ins,
const NameVarMap<egr::EagerVariable>& outs,
const framework::AttributeMap& attrs,
const framework::AttributeMap& default_attrs,
const phi::Place& place);
bool HasVoidFunctionPostHook() const {
return !void_function_post_hooks_.empty();
}
void AddVoidFunctionPostHook(std::shared_ptr<std::function<void()>>&& hook) {
void_function_post_hooks_.emplace_back(std::move(hook));
}
const std::vector<std::shared_ptr<std::function<void()>>>&
GetVoidFunctionPostHooks() const {
return void_function_post_hooks_;
}
private:
static const std::string& UnknownOpType() {
static std::string kUnknownOpType{"unknown"};
return kUnknownOpType;
}
private:
NameVarMap<VariableWrapper> ins_;
NameVarMap<VariableWrapper> outs_;
framework::AttributeMap attrs_;
const framework::AttributeMap* default_attrs_ = nullptr;
std::unique_ptr<framework::OperatorBase> op_;
phi::Place place_;
size_t id_{-1UL};
// In order to reduce the compatibility phase
// performance overhead, temporarily cache KernelContext
static phi::KernelContext phi_kernel_context_;
std::vector<std::shared_ptr<std::function<void()>>> void_function_post_hooks_;
};
class GradOpNode {
public:
GradOpNode() = default;
void reserve(size_t size) { ops_.reserve(size); }
size_t size() const { return ops_.size(); }
bool empty() const { return ops_.empty(); }
void clear() { ops_.clear(); }
void pop_back() { ops_.pop_back(); }
template <typename... ARGS>
OpBase& emplace_back(ARGS&&... args) { // NOLINT
ops_.emplace_back(std::forward<ARGS>(args)...);
return ops_.back();
}
const OpBase& back() const { return ops_.back(); }
OpBase& back() { return ops_.back(); }
OpBase& operator[](size_t idx) { return ops_[idx]; }
const OpBase& operator[](size_t idx) const { return ops_[idx]; }
/* Iterator related */
using Iterator = std::vector<OpBase>::iterator;
using ConstIterator = std::vector<OpBase>::const_iterator;
Iterator begin() { return ops_.begin(); }
Iterator end() { return ops_.end(); }
ConstIterator begin() const { return ops_.begin(); }
ConstIterator end() const { return ops_.end(); }
void InsertGradPendingNode(const std::shared_ptr<GradOpNode>& node) {
if (node && std::find(grad_pending_nodes_.begin(),
grad_pending_nodes_.end(),
node) == grad_pending_nodes_.end()) {
grad_pending_nodes_.emplace_back(node);
}
}
void SetInplaceGradNameMap(
const std::map<std::string, std::string>& inplace_input_map) {
for (auto& pair : inplace_input_map) {
VLOG(10) << "Set mapping relationship ("
<< framework::GradVarName(pair.first) << ", "
<< framework::GradVarName(pair.second)
<< ") for Inplace grad node.";
inplace_grad_name_map_[framework::GradVarName(pair.first)] =
framework::GradVarName(pair.second);
}
}
const std::map<std::string, std::string>& InplaceGradNameMap() const {
return inplace_grad_name_map_;
}
const std::vector<std::shared_ptr<GradOpNode>>& GradPendingNodes() const {
return grad_pending_nodes_;
}
private:
DISABLE_COPY_AND_ASSIGN(GradOpNode);
private:
std::vector<OpBase> ops_;
std::vector<std::shared_ptr<GradOpNode>> grad_pending_nodes_;
// Mapping relationship between grad output and grad input of the grad node of
// Inplace op.
std::map<std::string, std::string> inplace_grad_name_map_;
};
} // namespace imperative
} // namespace paddle