wehub-resource-sync
259 lines
9.9 KiB
C++
259 lines
9.9 KiB
C++
// Copyright (c) 2024 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 "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/kernels/funcs/compound_functors.h"
|
|
#include "paddle/phi/kernels/funcs/elementwise/elementwise_op_function.h"
|
|
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
|
|
#include "paddle/phi/kernels/funcs/functors.h"
|
|
#include "paddle/phi/kernels/funcs/fused_elemwise_activation_functor.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void FusedElemwiseActivationKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const DenseTensor &y,
|
|
const std::vector<std::string> &functor_list,
|
|
int axis,
|
|
float scale,
|
|
bool save_intermediate_out,
|
|
DenseTensor *out,
|
|
DenseTensor *intermediate_out) {
|
|
auto &in_x = GET_DATA_SAFELY(&x, "Input", "X", "FusedElemwiseActivation");
|
|
auto &in_y = GET_DATA_SAFELY(&y, "Input", "Y", "FusedElemwiseActivation");
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
out != nullptr,
|
|
true,
|
|
common::errors::InvalidArgument("The output(Out) should not be empty"));
|
|
auto output = out;
|
|
|
|
std::vector<DenseTensor *> outputs;
|
|
outputs.emplace_back(output);
|
|
|
|
if (save_intermediate_out) {
|
|
PADDLE_ENFORCE_EQ(intermediate_out != nullptr,
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"The save_intermediate_out is enable, so the "
|
|
"IntermediateOut should not be empty."));
|
|
|
|
outputs.emplace_back(intermediate_out);
|
|
} else {
|
|
outputs.emplace_back(nullptr);
|
|
}
|
|
|
|
funcs::RunFunctors<Context, T>(dev_ctx,
|
|
in_x,
|
|
in_y,
|
|
&outputs,
|
|
functor_list,
|
|
scale,
|
|
axis,
|
|
save_intermediate_out);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FusedElemwiseActivationGradKernel(
|
|
const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const DenseTensor &y,
|
|
const DenseTensor &out,
|
|
const DenseTensor &intermediate_out,
|
|
const DenseTensor &out_grad,
|
|
const std::vector<std::string> &functor_list,
|
|
int axis,
|
|
float scale,
|
|
bool save_intermediate_out,
|
|
DenseTensor *x_grad,
|
|
DenseTensor *y_grad) {
|
|
auto *in_y = &y;
|
|
PADDLE_ENFORCE_NE(
|
|
in_y,
|
|
nullptr,
|
|
common::errors::InvalidArgument("Input(Y) should not be nullptr."));
|
|
DenseTensor *in_out = const_cast<DenseTensor *>(&out);
|
|
|
|
auto in_out_grad = &out_grad;
|
|
PADDLE_ENFORCE_NE(in_out_grad,
|
|
nullptr,
|
|
common::errors::InvalidArgument(
|
|
"Input(Out@GRAD) should not be nullptr."));
|
|
|
|
std::vector<std::string> functor_list_new = functor_list;
|
|
size_t sz = functor_list_new[0].size();
|
|
int start = sz < 5 ? 0 : (sz - 5);
|
|
if (functor_list_new[0].substr(start, 5) != "_grad") {
|
|
functor_list_new[0] += "_grad";
|
|
}
|
|
sz = functor_list_new[1].size();
|
|
start = sz < 5 ? 0 : (sz - 5);
|
|
if (functor_list_new[1].substr(start, 5) != "_grad") {
|
|
functor_list_new[1] += "_grad";
|
|
}
|
|
|
|
DenseTensor *in_x = const_cast<DenseTensor *>(&x);
|
|
DenseTensor *d_intermediate_out =
|
|
nullptr; // intermediate_out_grad is not supported in ops.yaml, so use
|
|
// nullptr
|
|
|
|
// Get intermediate_out
|
|
DenseTensor *in_intermediate_out = nullptr;
|
|
if (save_intermediate_out) {
|
|
// if save_intermediate_out is true, for Unary(Binary(x, y)) and
|
|
// Binary(x, Unary(y)), the Binary(x, y) and Unary(y) not need to
|
|
// recompute.
|
|
in_intermediate_out = const_cast<DenseTensor *>(&intermediate_out);
|
|
PADDLE_ENFORCE_NE(in_intermediate_out,
|
|
nullptr,
|
|
common::errors::InvalidArgument(
|
|
"The option of 'save_intermediate_out' is opened,"
|
|
" so the number of 'Out' should be two."));
|
|
} else {
|
|
if (!funcs::InputXCanBeAbsent(functor_list_new)) {
|
|
PADDLE_ENFORCE_NE(
|
|
in_x,
|
|
nullptr,
|
|
common::errors::InvalidArgument("Input(X) should not be null."));
|
|
}
|
|
}
|
|
|
|
// Get in_x
|
|
if (x.initialized()) {
|
|
PADDLE_ENFORCE_NE(
|
|
in_x,
|
|
nullptr,
|
|
common::errors::InvalidArgument("Input(X) should not be null."));
|
|
} else {
|
|
// If functor_list contains elementwise_add, the backward doesn't use
|
|
// in_x, in_y and in_out.
|
|
PADDLE_ENFORCE_EQ(funcs::InputXCanBeAbsent(functor_list_new),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Only when the compoundfunctor contains "
|
|
"elementwise_add_grad, the 'X' could be absent."));
|
|
in_x = const_cast<DenseTensor *>(in_out_grad);
|
|
}
|
|
|
|
// Get in_Out
|
|
if (out.initialized()) {
|
|
PADDLE_ENFORCE_NE(
|
|
in_out,
|
|
nullptr,
|
|
common::errors::InvalidArgument("Input(X) should not be null."));
|
|
} else {
|
|
// If functor_list contains elementwise_add, the backward doesn't use
|
|
// in_x, in_y and in_out.
|
|
PADDLE_ENFORCE_EQ(funcs::InputXCanBeAbsent(functor_list_new),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Only when the compoundfunctor contains "
|
|
"elementwise_add_grad, the 'X' could be absent."));
|
|
in_out = const_cast<DenseTensor *>(in_out_grad);
|
|
}
|
|
|
|
bool has_in_place = funcs::HasInPlaceUnary(functor_list_new);
|
|
if (has_in_place) {
|
|
funcs::RunGradFunctors<Context, T, true /*InPlace*/>(dev_ctx,
|
|
in_x,
|
|
in_y,
|
|
in_out,
|
|
in_intermediate_out,
|
|
in_out_grad,
|
|
x_grad,
|
|
y_grad,
|
|
d_intermediate_out,
|
|
functor_list_new,
|
|
scale,
|
|
axis);
|
|
} else {
|
|
funcs::RunGradFunctors<Context, T, false /*InPlace*/>(dev_ctx,
|
|
in_x,
|
|
in_y,
|
|
in_out,
|
|
in_intermediate_out,
|
|
in_out_grad,
|
|
x_grad,
|
|
y_grad,
|
|
d_intermediate_out,
|
|
functor_list_new,
|
|
scale,
|
|
axis);
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FusedElemwiseAddActivationKernel(
|
|
const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const DenseTensor &y,
|
|
const std::vector<std::string> &functor_list,
|
|
int axis,
|
|
float scale,
|
|
bool save_intermediate_out,
|
|
DenseTensor *out,
|
|
DenseTensor *intermediate_out) {
|
|
FusedElemwiseActivationKernel<T, Context>(dev_ctx,
|
|
x,
|
|
y,
|
|
functor_list,
|
|
axis,
|
|
scale,
|
|
save_intermediate_out,
|
|
out,
|
|
intermediate_out);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FusedElemwiseAddActivationGradKernel(
|
|
const Context &dev_ctx,
|
|
const optional<DenseTensor> &x,
|
|
const DenseTensor &y,
|
|
const DenseTensor &out,
|
|
const optional<DenseTensor> &intermediate_out,
|
|
const DenseTensor &out_grad,
|
|
const std::vector<std::string> &functor_list,
|
|
int axis,
|
|
float scale,
|
|
bool save_intermediate_out,
|
|
DenseTensor *x_grad,
|
|
DenseTensor *y_grad) {
|
|
DenseTensor tmp_x;
|
|
DenseTensor tmp_i;
|
|
if (x) {
|
|
tmp_x = x.get();
|
|
}
|
|
if (intermediate_out) {
|
|
tmp_i = intermediate_out.get();
|
|
}
|
|
FusedElemwiseActivationGradKernel<T, Context>(dev_ctx,
|
|
tmp_x,
|
|
y,
|
|
out,
|
|
tmp_i,
|
|
out_grad,
|
|
functor_list,
|
|
axis,
|
|
scale,
|
|
save_intermediate_out,
|
|
x_grad,
|
|
y_grad);
|
|
}
|
|
|
|
} // namespace phi
|