Skip to content

[CustomOP Optional] CustomOP supports optional vector<Tensor> input #51973

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 23, 2023
Merged

[CustomOP Optional] CustomOP supports optional vector<Tensor> input #51973

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion paddle/fluid/framework/custom_operator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -174,14 +174,20 @@ static void RunKernelFunc(
custom_t.set_impl(std::make_shared<phi::DenseTensor>(*x));
custom_vec_in.emplace_back(custom_t);
}
} else { // optional inputs, `custom_vec_in` is empty
} else { // optional inputs.
PADDLE_ENFORCE(
detail::IsOptionalVar(in_name),
phi::errors::NotFound("Your custom operator's KernelFunc cannot "
"find input parameter `%s`",
in_name));
VLOG(3) << "Custom Operator: KernelFunc's vector input " << in_name
<< " is optional dtype with None input";
// NOTE(HongyuJia): In dygraph mode, we can not distinguish Tensor and
// vector<Tensor> when user inputs None, so dygraph mode appends one
// un-initialized Tensor to CustomOpKernelContext. To be compatible with
// dygraph mode, `custom_vec_in` also emplace_back one un-initialized
// tensor here.
custom_vec_in.emplace_back(paddle::Tensor());
}
kernel_ctx.EmplaceBackInputs(std::move(custom_vec_in));
} else { // inputs Tensor
Expand Down
4 changes: 3 additions & 1 deletion paddle/fluid/pybind/pybind.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1060,7 +1060,9 @@ PYBIND11_MODULE(libpaddle, m) {
if (PyList_Check(obj) || PyTuple_Check(obj)) {
self.EmplaceBackInputs(
std::move(CastPyArg2VectorOfTensor(obj, 1)));
} else if (obj == Py_None) { // check optional Tensor
} else if (obj == Py_None) {
// Check optional Tensor, use one un-initialized tensor to
// indicate both Tensor and vector<Tensor> inputs
self.EmplaceBackInput(std::move(paddle::Tensor()));
} else {
self.EmplaceBackInput(std::move(CastPyArg2Tensor(obj, 1)));
Expand Down
71 changes: 69 additions & 2 deletions paddle/phi/api/ext/op_meta_info.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,6 +241,26 @@ struct KernelFuncImpl<Return (*)(Args...), impl_fn> {
}
};

template <typename... Tail>
struct ComputeCallHelper<const paddle::optional<std::vector<paddle::Tensor>>&,
Tail...> {
template <int in_idx, int attr_idx, int out_idx, typename... PreviousArgs>
static void Compute(CustomOpKernelContext* ctx, PreviousArgs&... pargs) {
auto& range = ctx->InputRangeAt(in_idx);
auto arg = ctx->InputsBetween(range.first, range.second);
if (arg.empty() || !arg[0].is_initialized()) {
ComputeCallHelper<Tail...>::
template Compute<in_idx + 1, attr_idx, out_idx>(
ctx, pargs..., paddle::none);
} else {
ComputeCallHelper<
Tail...>::template Compute<in_idx + 1, attr_idx, out_idx>(ctx,
pargs...,
arg);
}
}
};

PD_SPECIALIZE_ComputeCallHelper(bool);
PD_SPECIALIZE_ComputeCallHelper(int);
PD_SPECIALIZE_ComputeCallHelper(float);
Expand Down Expand Up @@ -486,6 +506,33 @@ struct InferShapeFuncImpl<Return (*)(Args...), impl_fn> {
}
};

template <typename... Tail>
struct InferShapeCallHelper<
const paddle::optional<std::vector<std::vector<int64_t>>>&,
Tail...> {
template <int in_idx,
int vec_in_idx,
int attr_idx,
typename... PreviousArgs>
static Return InferShape(
const std::vector<std::vector<int64_t>>& input_shapes,
const std::vector<std::vector<std::vector<int64_t>>>& vec_input_shapes,
const std::vector<paddle::any>& attrs,
const PreviousArgs&... pargs) {
const std::vector<std::vector<int64_t>>& arg =
vec_input_shapes[vec_in_idx];
if (arg.empty()) {
return InferShapeCallHelper<Tail...>::
template InferShape<in_idx, vec_in_idx + 1, attr_idx>(
input_shapes, vec_input_shapes, attrs, pargs..., paddle::none);
} else {
return InferShapeCallHelper<Tail...>::
template InferShape<in_idx, vec_in_idx + 1, attr_idx>(
input_shapes, vec_input_shapes, attrs, pargs..., arg);
}
}
};

// NOTE(chenweihang): Used to be compatible with the 2.0.1 released
// interface, and will be deprecated in the future
PD_SPECIALIZE_InferShapeCallHelper_FOR_SHAPE(std::vector<int64_t>);
Expand Down Expand Up @@ -593,8 +640,7 @@ struct InferDtypeFuncImpl<Return (*)(Args...), impl_fn> {
PD_SPECIALIZE_InferDtypeCallHelper_FOR_DTYPES(const std::vector<DataType>&);

template <typename... Tail>
struct InferDtypeCallHelper<const paddle::optional<paddle::DataType>&,
Tail...> {
struct InferDtypeCallHelper<const paddle::optional<DataType>&, Tail...> {
template <int in_idx, int vec_in_idx, typename... PreviousArgs>
static Return InferDtype(
const std::vector<DataType>& input_dtypes,
Expand All @@ -613,6 +659,27 @@ struct InferDtypeFuncImpl<Return (*)(Args...), impl_fn> {
}
};

template <typename... Tail>
struct InferDtypeCallHelper<const paddle::optional<std::vector<DataType>>&,
Tail...> {
template <int in_idx, int vec_in_idx, typename... PreviousArgs>
static Return InferDtype(
const std::vector<DataType>& input_dtypes,
const std::vector<std::vector<DataType>>& vec_input_dtypes,
const PreviousArgs&... pargs) {
const std::vector<DataType>& arg = vec_input_dtypes[vec_in_idx];
if (arg.empty()) {
return InferDtypeCallHelper<Tail...>::
template InferDtype<in_idx, vec_in_idx + 1>(
input_dtypes, vec_input_dtypes, pargs..., paddle::none);
} else {
return InferDtypeCallHelper<Tail...>::
template InferDtype<in_idx, vec_in_idx + 1>(
input_dtypes, vec_input_dtypes, pargs..., arg);
}
}
};

// NOTE(chenweihang): Used to be compatible with the 2.0.1 released
// interface, and will be deprecated in the future
PD_SPECIALIZE_InferDtypeCallHelper_TO_DTYPE(DataType);
Expand Down
129 changes: 105 additions & 24 deletions python/paddle/fluid/tests/custom_op/custom_optional.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,21 +19,19 @@
#include "paddle/extension.h"

template <typename data_t>
void add_forward_kernel(const data_t* x_data,
const data_t* y_data,
data_t* out_data,
int64_t numel) {
void add_one_pointer(const data_t* x_data, data_t* out_data, int64_t numel) {
for (size_t i = 0; i < numel; ++i) {
out_data[i] = x_data[i] + y_data[i];
out_data[i] += x_data[i];
}
}

template <typename data_t>
void add_backward_kernel(data_t* x_grad_data,
const data_t* out_grad_data,
int64_t numel) {
void add_two_pointers(const data_t* x_data,
const data_t* y_data,
data_t* out_data,
int64_t numel) {
for (size_t i = 0; i < numel; ++i) {
x_grad_data[i] += out_grad_data[i];
out_data[i] = x_data[i] + y_data[i];
}
}

Expand All @@ -53,12 +51,12 @@ std::vector<paddle::Tensor> AddForward(
PD_DISPATCH_FLOATING_TYPES(
x.type(), "AddForward", ([&] {
if (y) {
add_forward_kernel<data_t>(x.data<data_t>(),
y->data<data_t>(),
out.data<data_t>(),
x.size());
add_two_pointers<data_t>(x.data<data_t>(),
y->data<data_t>(),
out.data<data_t>(),
x.size());
} else {
add_forward_kernel<data_t>(
add_two_pointers<data_t>(
x.data<data_t>(), x.data<data_t>(), out.data<data_t>(), x.size());
}
}));
Expand All @@ -69,7 +67,6 @@ std::vector<paddle::DataType> AddInferDtype(
const paddle::DataType& x_dtype,
const paddle::optional<paddle::DataType>& y_dtype) {
if (y_dtype) {
std::cout << "DEBUG AddInferDtype" << *y_dtype << std::endl;
return {*y_dtype};
}
return {x_dtype};
Expand Down Expand Up @@ -98,18 +95,14 @@ std::vector<paddle::Tensor> AddBackward(
PD_CHECK(x.place() == paddle::PlaceType::kCPU, "x must be a CPU Tensor.");

paddle::Tensor x_grad = paddle::zeros(x.shape(), x.dtype(), x.place());
paddle::Tensor y_grad = paddle::zeros(x.shape(), x.dtype(), x.place());

PD_DISPATCH_FLOATING_TYPES(
out_grad.type(), "AddBackward", ([&] {
add_backward_kernel<data_t>(
x_grad.data<data_t>(), out_grad.data<data_t>(), out_grad.size());
if (y) {
add_backward_kernel<data_t>(
y_grad.data<data_t>(), out_grad.data<data_t>(), out_grad.size());
} else {
add_backward_kernel<data_t>(
x_grad.data<data_t>(), out_grad.data<data_t>(), out_grad.size());
add_one_pointer<data_t>(
out_grad.data<data_t>(), x_grad.data<data_t>(), out_grad.size());
if (!y) {
add_one_pointer<data_t>(
out_grad.data<data_t>(), x_grad.data<data_t>(), out_grad.size());
}
}));

Expand All @@ -127,3 +120,91 @@ PD_BUILD_GRAD_OP(custom_add)
.Inputs({"X", paddle::Optional("Y"), paddle::Grad("Out")})
.Outputs({paddle::Grad("X")})
.SetKernelFn(PD_KERNEL(AddBackward));

/*
if (y) {
out = x + y[0] + y[1] + ...;
} else {
out = x + x;
}
*/
std::vector<paddle::Tensor> AddVectorForward(
const paddle::Tensor& x,
const paddle::optional<std::vector<paddle::Tensor>>& y) { // NOLINT
PD_CHECK(x.place() == paddle::PlaceType::kCPU, "x must be a CPU Tensor.");
paddle::Tensor out = paddle::zeros(x.shape(), x.dtype(), x.place());

PD_DISPATCH_FLOATING_TYPES(
x.type(), "AddVectorForward", ([&] {
if (y) {
add_one_pointer<data_t>(
x.data<data_t>(), out.data<data_t>(), out.size());
for (size_t i = 0; i < y->size(); ++i) {
add_one_pointer<data_t>(
y->at(i).data<data_t>(), out.data<data_t>(), out.size());
}
} else {
add_two_pointers<data_t>(
x.data<data_t>(), x.data<data_t>(), out.data<data_t>(), x.size());
}
}));
return {out};
}

std::vector<paddle::DataType> AddVectorInferDtype(
const paddle::DataType& x_dtype,
const paddle::optional<std::vector<paddle::DataType>>& y_dtype) {
if (y_dtype) {
return {y_dtype->at(0)};
}
return {x_dtype};
}

std::vector<std::vector<int64_t>> AddVectorInferShape(
const std::vector<int64_t>& x_shape,
const paddle::optional<std::vector<std::vector<int64_t>>>& y_shape) {
if (y_shape) {
return {y_shape->at(0)};
}
return {x_shape};
}

/*
if (y) {
x_grad = out_grad;
} else {
x_grad = out_grad + out_grad;
}
*/
std::vector<paddle::Tensor> AddVectorBackward(
const paddle::Tensor& x,
const paddle::optional<std::vector<paddle::Tensor>>& y,
const paddle::Tensor& out_grad) { // NOLINT
PD_CHECK(x.place() == paddle::PlaceType::kCPU, "x must be a CPU Tensor.");

paddle::Tensor x_grad = paddle::zeros(x.shape(), x.dtype(), x.place());

PD_DISPATCH_FLOATING_TYPES(
out_grad.type(), "AddVectorBackward", ([&] {
add_one_pointer<data_t>(
out_grad.data<data_t>(), x_grad.data<data_t>(), out_grad.size());
if (!y) {
add_one_pointer<data_t>(
out_grad.data<data_t>(), x_grad.data<data_t>(), out_grad.size());
}
}));

return {x_grad};
}

PD_BUILD_OP(custom_add_vec)
.Inputs({"X", paddle::Optional(paddle::Vec("Y"))})
.Outputs({"Out"})
.SetKernelFn(PD_KERNEL(AddVectorForward))
.SetInferShapeFn(PD_INFER_SHAPE(AddVectorInferShape))
.SetInferDtypeFn(PD_INFER_DTYPE(AddVectorInferDtype));

PD_BUILD_GRAD_OP(custom_add_vec)
.Inputs({"X", paddle::Optional(paddle::Vec("Y")), paddle::Grad("Out")})
.Outputs({paddle::Grad("X")})
.SetKernelFn(PD_KERNEL(AddVectorBackward));
Loading