Files
2026-07-13 12:40:42 +08:00

296 lines
12 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.
#include <string>
#include "paddle/fluid/framework/convert_utils.h"
#include "paddle/fluid/framework/infershape_utils.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/op_version_registry.h"
#include "paddle/fluid/framework/tensor_util.h"
#include "paddle/phi/core/infermeta_utils.h"
#include "paddle/phi/infermeta/unary.h"
namespace paddle::framework {
class InferShapeContext;
class OpDesc;
template <typename T>
class EmptyGradOpMaker;
} // namespace paddle::framework
namespace paddle::imperative {
class OpBase;
} // namespace paddle::imperative
namespace paddle::operators {
class SetValue : public framework::OperatorWithKernel {
public:
SetValue(const std::string &type,
const framework::VariableNameMap &inputs,
const framework::VariableNameMap &outputs,
const framework::AttributeMap &attrs)
: OperatorWithKernel(type, inputs, outputs, attrs) {}
protected:
phi::KernelKey GetExpectedKernelType(
const framework::ExecutionContext &ctx) const override {
return phi::KernelKey(OperatorWithKernel::IndicateVarDataType(ctx, "Input"),
ctx.GetPlace());
}
phi::KernelKey GetKernelTypeForVar(
const std::string &var_name,
const DenseTensor &tensor,
const phi::KernelKey &expected_kernel_type) const override {
if (var_name == "StartsTensorList" || var_name == "EndsTensorList" ||
var_name == "StepsTensorList") {
return phi::KernelKey(phi::Backend::ALL_BACKEND,
expected_kernel_type.layout(),
expected_kernel_type.dtype());
}
return phi::KernelKey(
tensor.place(), tensor.layout(), expected_kernel_type.dtype());
}
};
class SetValueMaker : public framework::OpProtoAndCheckerMaker {
public:
void Make() override {
// Input
AddInput("Input", "(phi::DenseTensor) Input tensor of set_value operator.");
AddInput("ValueTensor",
"(phi::DenseTensor) Value tensor of set_value operator.")
.AsDispensable();
AddInput("StartsTensorList",
"(vector<phi::DenseTensor<int32>>, optional) If provided, "
"set_value will "
"use this. The shape of the tensor in vector must be [1]."
"It has higher priority compare with attr(starts).")
.AsDuplicable()
.AsDispensable();
AddInput("EndsTensorList",
"(vector<phi::DenseTensor<int32>>, optional) If provided, "
"set_value will "
"use this. The shape of the tensor in vector must BE [1]."
"It has higher priority compare with attr(ends).")
.AsDuplicable()
.AsDispensable();
AddInput("StepsTensorList",
"(vector<phi::DenseTensor<int32>>, optional) If provided, "
"set_value will "
"use this. The shape of the tensor in vector must BE [1]."
"It has higher priority compare with attr(steps).")
.AsDuplicable()
.AsDispensable();
// Output
AddOutput("Out",
"(phi::DenseTensor) Output tensor of set_value operator. The "
"output is the "
"same DenseTensor as input");
// Attr
AddAttr<int>("dtype", "data type of input.")
.InEnum({framework::proto::VarType::BOOL,
framework::proto::VarType::INT32,
framework::proto::VarType::INT64,
framework::proto::VarType::FP32,
framework::proto::VarType::FP64,
framework::proto::VarType::FP16,
framework::proto::VarType::BF16,
framework::proto::VarType::COMPLEX64,
framework::proto::VarType::COMPLEX128})
.SetDefault(framework::proto::VarType::FP32);
AddAttr<std::vector<int64_t>>(
"axes", "(list<int64_t>) Axes that `starts` and `ends` apply to.");
AddAttr<std::vector<int64_t>>(
"starts",
"(list<int64_t>) Starting indices of corresponding axis in `axes`.")
.SetDefault({});
AddAttr<std::vector<int64_t>>(
"ends",
"(list<int64_t>) Ending indices of corresponding axis in `axes`.")
.SetDefault({});
AddAttr<std::vector<int64_t>>(
"steps", "(list<int64_t>) Stride step from the start to the end.")
.SetDefault({});
AddAttr<std::vector<int64_t>>("decrease_axes",
"(list<int>) The axes to decrease.")
.SetDefault({});
AddAttr<std::vector<int64_t>>("none_axes", "(list<int>) The axes to none.")
.SetDefault({});
AddAttr<std::vector<paddle::experimental::Scalar>>("values", "values")
.SetDefault({});
AddAttr<std::vector<int64_t>>("shape", "(vector<int64_t>) Shape of values.")
.SetDefault({});
AddComment(R"DOC(SetValue operator.
Assignment to a DenseTensor in static graph mode.
)DOC");
}
};
template <typename T>
class SetValueGradMaker : public framework::SingleGradOpMaker<T> {
public:
using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
protected:
void Apply(GradOpPtr<T> op) const override {
op->SetType("set_value_grad");
op->SetInput("ValueTensor", this->Input("ValueTensor"));
op->SetOutput(framework::GradVarName("ValueTensor"),
this->InputGrad("ValueTensor"));
op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
if (this->HasInput("StartsTensorList")) {
op->SetInput("StartsTensorList", this->Input("StartsTensorList"));
}
if (this->HasInput("EndsTensorList")) {
op->SetInput("EndsTensorList", this->Input("EndsTensorList"));
}
if (this->HasInput("StepsTensorList")) {
op->SetInput("StepsTensorList", this->Input("StepsTensorList"));
}
op->SetAttrMap(this->Attrs());
op->SetOutput(framework::GradVarName("Input"), this->InputGrad("Input"));
}
};
class SetValueGrad : public framework::OperatorWithKernel {
public:
using framework::OperatorWithKernel::OperatorWithKernel;
void InferShape(framework::InferShapeContext *ctx) const override {
OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Out")),
"Input",
framework::GradVarName("Out"),
"set_value_grad");
auto in_dims = ctx->GetInputDim(framework::GradVarName("Out"));
PADDLE_ENFORCE_LT(
in_dims.size(),
7,
common::errors::InvalidArgument(
"The dimension of set_value_grad operator's input should be less "
"than 7, but received dimension is %d.",
in_dims.size()));
if (ctx->HasOutput(framework::GradVarName("ValueTensor"))) {
ctx->ShareDim("ValueTensor",
/*->*/ framework::GradVarName("ValueTensor"));
ctx->ShareLoD("ValueTensor",
/*->*/ framework::GradVarName("ValueTensor"));
}
}
protected:
phi::KernelKey GetExpectedKernelType(
const framework::ExecutionContext &ctx) const override {
auto in_tensor = ctx.Input<DenseTensor>(framework::GradVarName("Out"));
return phi::KernelKey(OperatorWithKernel::IndicateVarDataType(
ctx, framework::GradVarName("Out")),
in_tensor->place());
}
phi::KernelKey GetKernelTypeForVar(
const std::string &var_name,
const DenseTensor &tensor,
const phi::KernelKey &expected_kernel_type) const override {
if (var_name == "StartsTensorList" || var_name == "EndsTensorList" ||
var_name == "StepsTensorList") {
return phi::KernelKey(phi::Backend::ALL_BACKEND,
expected_kernel_type.layout(),
expected_kernel_type.dtype());
}
return phi::KernelKey(
tensor.place(), tensor.layout(), expected_kernel_type.dtype());
}
};
DECLARE_INPLACE_OP_INFERER(SetValueOpInplaceInferer, {"Input", "Out"});
} // namespace paddle::operators
namespace ops = paddle::operators;
DECLARE_INFER_SHAPE_FUNCTOR(set_value,
SetValueInferShapeFunctor,
PD_INFER_META(phi::SetValueInferMeta));
REGISTER_OPERATOR(set_value,
ops::SetValue,
ops::SetValueMaker,
ops::SetValueGradMaker<paddle::framework::OpDesc>,
ops::SetValueGradMaker<paddle::imperative::OpBase>,
ops::SetValueOpInplaceInferer,
SetValueInferShapeFunctor);
REGISTER_OPERATOR(set_value_grad, ops::SetValueGrad);
REGISTER_OP_VERSION(set_value)
.AddCheckpoint(
R"ROC(
Upgrade set_value, add 3 inputs [StartsTensorList, EndsTensorList, StepsTensorList] and 1 attribute [steps].
)ROC",
paddle::framework::compatible::OpVersionDesc()
.NewInput("StartsTensorList",
"If provided, set_value will use this.The shape of the "
"tensor in vector must be [1]. It has higher priority "
"compare with attr(starts).")
.NewInput("EndsTensorList",
"If provided, set_value will use this.The shape of the "
"tensor in vector must be [1]. It has higher priority "
"compare with attr(ends).")
.NewInput("StepsTensorList",
"If provided, set_value will use this.The shape of the "
"tensor in vector must be [1]. It has higher priority "
"compare with attr(steps).")
.ModifyAttr("starts",
"Starting indices of corresponding axis in `axes`.",
std::vector<int64_t>{})
.ModifyAttr("ends",
"Ending indices of corresponding axis in `axes`.",
std::vector<int64_t>{})
.NewAttr("steps",
"Stride step from the start to the end.",
std::vector<int64_t>{}))
.AddCheckpoint(
R"ROC(
Upgrade set_value, add 1 attribute [decrease_axes].
)ROC",
paddle::framework::compatible::OpVersionDesc().NewAttr(
"decrease_axes", "The axes to decrease.", std::vector<int64_t>{}))
.AddCheckpoint(
R"ROC(
Upgrade set_value, add 1 attribute [none_axes].
)ROC",
paddle::framework::compatible::OpVersionDesc().NewAttr(
"none_axes", "The axes with none index.", std::vector<int64_t>{}))
.AddCheckpoint(
R"ROC(Upgrade set_value to support generic Scalars as value and remove plain values, so as to support complex types.)ROC",
paddle::framework::compatible::OpVersionDesc()
.NewAttr("values",
"values",
std::vector<paddle::experimental::Scalar>())
.DeleteAttr("bool_values", "remove plain attributes")
.DeleteAttr("fp32_values", "remove plain attributes")
.DeleteAttr("int32_values", "remove plain attributes")
.DeleteAttr("int64_values", "remove plain attributes")
.DeleteAttr("fp64_values", "remove plain attributes"));