wehub-resource-sync
1525 lines
62 KiB
C++
1525 lines
62 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/phi/api/lib/api_custom_impl.h"
|
|
|
|
#include "glog/logging.h"
|
|
#include "paddle/common/flags.h"
|
|
#include "paddle/phi/api/lib/api_gen_utils.h"
|
|
#include "paddle/phi/api/lib/data_transform.h"
|
|
#include "paddle/phi/api/lib/kernel_dispatch.h"
|
|
#include "paddle/phi/api/lib/tensor_copy.h"
|
|
#include "paddle/phi/common/type_traits.h"
|
|
#include "paddle/phi/core/compat/convert_utils.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/meta_tensor.h"
|
|
#include "paddle/phi/infermeta/backward.h"
|
|
#include "paddle/phi/infermeta/binary.h"
|
|
#include "paddle/phi/infermeta/fusion.h"
|
|
#include "paddle/phi/infermeta/multiary.h"
|
|
#include "paddle/phi/infermeta/nullary.h"
|
|
#include "paddle/phi/infermeta/unary.h"
|
|
|
|
#include "paddle/phi/api/profiler/event_tracing.h"
|
|
#include "paddle/phi/api/profiler/supplement_tracing.h"
|
|
#ifdef PADDLE_WITH_DISTRIBUTE
|
|
#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
|
|
#include "paddle/phi/infermeta/spmd_rules/rules.h"
|
|
#endif
|
|
|
|
COMMON_DECLARE_int32(low_precision_op_list);
|
|
COMMON_DECLARE_bool(benchmark);
|
|
|
|
namespace paddle::experimental {
|
|
|
|
////////////////// Forward api impls //////////////////////
|
|
|
|
Tensor add_n_impl(const std::vector<Tensor>& x) {
|
|
Backend kernel_backend = Backend::UNDEFINED;
|
|
DataLayout kernel_layout = DataLayout::UNDEFINED;
|
|
DataType kernel_data_type = DataType::UNDEFINED;
|
|
|
|
if (kernel_backend == Backend::UNDEFINED ||
|
|
kernel_layout == DataLayout::UNDEFINED ||
|
|
kernel_data_type == DataType::UNDEFINED) {
|
|
auto kernel_key_set = ParseKernelKeyByInputArgs(x);
|
|
auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey();
|
|
if (kernel_backend == Backend::UNDEFINED) {
|
|
kernel_backend = kernel_key.backend();
|
|
}
|
|
if (kernel_layout == DataLayout::UNDEFINED) {
|
|
kernel_layout = kernel_key.layout();
|
|
}
|
|
if (kernel_data_type == DataType::UNDEFINED) {
|
|
kernel_data_type = kernel_key.dtype();
|
|
}
|
|
}
|
|
|
|
bool is_sr_kernel = true;
|
|
for (auto& input : x) {
|
|
if (phi::DenseTensor::classof(input.impl().get()) ||
|
|
phi::distributed::DistTensor::classof(input.impl().get())) {
|
|
is_sr_kernel = false;
|
|
break;
|
|
}
|
|
}
|
|
|
|
const std::string kernel_name = (is_sr_kernel ? "add_n_sr" : "add_n");
|
|
|
|
VLOG(6) << "add_n API kernel key: [" << kernel_backend << ", "
|
|
<< kernel_layout << ", " << kernel_data_type << "]";
|
|
auto kernel_result = phi::KernelFactory::Instance().SelectKernelOrThrowError(
|
|
kernel_name, {kernel_backend, kernel_layout, kernel_data_type});
|
|
const auto& kernel = kernel_result.kernel;
|
|
VLOG(6) << kernel_name << " kernel: " << kernel;
|
|
auto* dev_ctx = GetDeviceContextByBackend(
|
|
kernel_result.has_fallback_cpu ? Backend::CPU : kernel_backend);
|
|
|
|
Tensor api_output;
|
|
|
|
if (is_sr_kernel) {
|
|
std::vector<const phi::SelectedRows*> input_x(x.size());
|
|
for (size_t i = 0; i < input_x.size(); ++i) {
|
|
input_x[i] = static_cast<phi::SelectedRows*>(x[i].impl().get());
|
|
}
|
|
auto x_meta_vec = MakeMetaTensor(input_x);
|
|
std::vector<const phi::MetaTensor*> x_metas(x_meta_vec.size());
|
|
for (size_t i = 0; i < x_meta_vec.size(); ++i) {
|
|
x_metas[i] = &x_meta_vec[i];
|
|
}
|
|
auto kernel_out = SetSelectedRowsKernelOutput(&api_output);
|
|
phi::MetaTensor meta_out(kernel_out);
|
|
phi::AddNInferMeta(x_metas, &meta_out);
|
|
|
|
using kernel_signature =
|
|
void (*)(const phi::DeviceContext&,
|
|
const std::vector<const phi::SelectedRows*>&,
|
|
phi::SelectedRows*);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
|
|
(*kernel_fn)(*dev_ctx, input_x, kernel_out);
|
|
} else {
|
|
#ifdef PADDLE_WITH_DISTRIBUTE
|
|
bool run_auto_parallel = AllInputsAreDistTensor(x);
|
|
bool rank_is_in_current_mesh = true;
|
|
if (run_auto_parallel) {
|
|
auto mesh =
|
|
std::static_pointer_cast<phi::distributed::DistTensor>(x[0].impl())
|
|
->dist_attr()
|
|
.process_mesh();
|
|
rank_is_in_current_mesh = phi::distributed::IsCurRankInMesh(mesh);
|
|
|
|
std::vector<const phi::TensorBase*> input_x(x.size());
|
|
for (size_t i = 0; i < input_x.size(); ++i) {
|
|
input_x[i] = x[i].impl().get();
|
|
}
|
|
|
|
auto meta_dist_input_x = MakeDistMetaTensor(input_x);
|
|
auto spmd_info = phi::distributed::VariadicReplicatedInferSpmdDynamic(
|
|
meta_dist_input_x);
|
|
|
|
auto dist_out = SetKernelDistOutput(&api_output);
|
|
auto dense_out = dist_out->unsafe_mutable_value();
|
|
if (!rank_is_in_current_mesh) {
|
|
*dense_out = phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(
|
|
nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}
|
|
|
|
phi::MetaTensor meta_dist_out(dist_out);
|
|
auto x_meta_vec = MakeMetaTensor(input_x);
|
|
std::vector<const phi::MetaTensor*> x_metas(x_meta_vec.size());
|
|
for (size_t i = 0; i < x_meta_vec.size(); ++i) {
|
|
x_metas[i] = &x_meta_vec[i];
|
|
}
|
|
phi::AddNInferMeta(x_metas, &meta_dist_out);
|
|
if (rank_is_in_current_mesh) {
|
|
auto dist_input_x =
|
|
ReshardApiInputToKernelInput(dev_ctx, x, spmd_info.first[0]);
|
|
dist_input_x = PrepareDataForDistTensor(
|
|
dist_input_x,
|
|
GetKernelInputArgDef(kernel.InputAt(0), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
std::vector<const phi::TensorBase*> input_x(dist_input_x.size());
|
|
for (size_t i = 0; i < dist_input_x.size(); ++i) {
|
|
input_x[i] = dist_input_x[i]->unsafe_mutable_value();
|
|
}
|
|
|
|
auto x_meta_vec = MakeMetaTensor(input_x);
|
|
std::vector<const phi::MetaTensor*> x_metas(x_meta_vec.size());
|
|
for (size_t i = 0; i < x_meta_vec.size(); ++i) {
|
|
x_metas[i] = &x_meta_vec[i];
|
|
}
|
|
phi::MetaTensor meta_dense_out(dense_out);
|
|
phi::AddNInferMeta(x_metas, &meta_dense_out);
|
|
|
|
using kernel_signature =
|
|
void (*)(const phi::DeviceContext&,
|
|
const std::vector<const phi::TensorBase*>&,
|
|
phi::DenseTensor*);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
(*kernel_fn)(*dev_ctx, input_x, dense_out);
|
|
}
|
|
PADDLE_ENFORCE_EQ(paddle::holds_alternative<
|
|
std::vector<phi::distributed::TensorDistAttr>>(
|
|
spmd_info.first[0]),
|
|
true,
|
|
common::errors::PreconditionNotMet(
|
|
"Arg must be a vector of TensorDistAttr"));
|
|
|
|
auto current_process_mesh =
|
|
paddle::get<1>(spmd_info.first[0]).at(0).process_mesh();
|
|
SetReplicatedDistAttrForOutput(dist_out, current_process_mesh);
|
|
return api_output;
|
|
}
|
|
#endif
|
|
std::vector<const phi::TensorBase*> input_x(x.size());
|
|
std::vector<std::shared_ptr<phi::DenseTensor>> temp_dense_tensors;
|
|
temp_dense_tensors.reserve(x.size());
|
|
for (size_t i = 0; i < input_x.size(); ++i) {
|
|
if (phi::DenseTensor::classof(x[i].impl().get())) {
|
|
temp_dense_tensors.push_back(
|
|
PrepareData(x[i], kernel.InputAt(0), {}, false));
|
|
input_x[i] = temp_dense_tensors.back().get();
|
|
} else {
|
|
input_x[i] = x[i].impl().get();
|
|
}
|
|
}
|
|
auto x_meta_vec = MakeMetaTensor(input_x);
|
|
std::vector<const phi::MetaTensor*> x_metas(x_meta_vec.size());
|
|
for (size_t i = 0; i < x_meta_vec.size(); ++i) {
|
|
x_metas[i] = &x_meta_vec[i];
|
|
}
|
|
auto kernel_out = SetKernelOutput(&api_output);
|
|
phi::MetaTensor meta_out(kernel_out);
|
|
|
|
std::vector<phi::MetaTensor*> output_metas_for_compact;
|
|
output_metas_for_compact.push_back(&meta_out);
|
|
phi::AddNInferMeta(x_metas, &meta_out);
|
|
CheckAndDoCompact(output_metas_for_compact, "add_n");
|
|
|
|
using kernel_signature =
|
|
void (*)(const phi::DeviceContext&,
|
|
const std::vector<const phi::TensorBase*>&,
|
|
phi::DenseTensor*);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
|
|
(*kernel_fn)(*dev_ctx, input_x, kernel_out);
|
|
if (kernel_result.has_fallback_cpu) {
|
|
TransDataBackend(kernel_out, kernel_backend, kernel_out);
|
|
}
|
|
}
|
|
|
|
return api_output;
|
|
}
|
|
|
|
Tensor copy_to_impl(const Tensor& x, Place place, bool blocking) {
|
|
Tensor out;
|
|
copy(x, place, blocking, &out);
|
|
out.set_name(x.name());
|
|
return out;
|
|
}
|
|
|
|
std::tuple<Tensor, Tensor> fused_gemm_epilogue_impl(
|
|
const Tensor& x,
|
|
const Tensor& y,
|
|
const Tensor& bias,
|
|
bool trans_x,
|
|
bool trans_y,
|
|
const std::string& activation) {
|
|
// Kernel Key Construction
|
|
Backend kernel_backend = Backend::UNDEFINED;
|
|
DataLayout kernel_layout = DataLayout::UNDEFINED;
|
|
DataType kernel_data_type = DataType::UNDEFINED;
|
|
if (kernel_backend == Backend::UNDEFINED ||
|
|
kernel_layout == DataLayout::UNDEFINED ||
|
|
kernel_data_type == DataType::UNDEFINED) {
|
|
auto kernel_key_set = ParseKernelKeyByInputArgs(x, y, bias);
|
|
auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey();
|
|
if (kernel_backend == Backend::UNDEFINED) {
|
|
kernel_backend = kernel_key.backend();
|
|
}
|
|
if (kernel_layout == DataLayout::UNDEFINED) {
|
|
kernel_layout = kernel_key.layout();
|
|
}
|
|
if (kernel_data_type == DataType::UNDEFINED) {
|
|
kernel_data_type = kernel_key.dtype();
|
|
}
|
|
}
|
|
#ifdef PADDLE_WITH_DISTRIBUTE
|
|
bool run_auto_parallel = AllInputsAreDistTensor(x, y, bias);
|
|
bool rank_is_in_current_mesh = true;
|
|
if (run_auto_parallel) {
|
|
auto mesh =
|
|
std::static_pointer_cast<phi::distributed::DistTensor>(bias.impl())
|
|
->dist_attr()
|
|
.process_mesh();
|
|
rank_is_in_current_mesh = phi::distributed::IsCurRankInMesh(mesh);
|
|
}
|
|
|
|
// Kernel Dispatch Body
|
|
// Auto Parallel condition
|
|
if (run_auto_parallel) {
|
|
// 1. InferSpmd (Infer DistAttr of Inputs&Outputs)
|
|
auto meta_dist_input_x = MakeDistMetaTensor(*x.impl());
|
|
auto meta_dist_input_y = MakeDistMetaTensor(*y.impl());
|
|
auto meta_dist_input_bias = MakeDistMetaTensor(*bias.impl());
|
|
auto spmd_info =
|
|
phi::distributed::FusedGemmEpilogueInferSpmd(meta_dist_input_x,
|
|
meta_dist_input_y,
|
|
meta_dist_input_bias,
|
|
trans_x,
|
|
trans_y,
|
|
activation);
|
|
DebugInfoForInferSpmd("fused_gemm_epilogue", spmd_info);
|
|
|
|
// 2. Create API Output & Prepare Dist and Dense Output
|
|
phi::DeviceContext* dev_ctx = nullptr;
|
|
std::tuple<Tensor, Tensor> api_output;
|
|
|
|
auto dist_out_0 =
|
|
SetKernelDistOutput(&std::get<0>(api_output), spmd_info.second[0]);
|
|
auto dense_out_0 =
|
|
dist_out_0 ? dist_out_0->unsafe_mutable_value() : nullptr;
|
|
if (!rank_is_in_current_mesh) {
|
|
*dense_out_0 =
|
|
phi::DenseTensor(std::make_shared<phi::Allocation>(
|
|
nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}
|
|
|
|
phi::distributed::DistTensor* dist_out_1 = nullptr;
|
|
if (activation != "none") {
|
|
dist_out_1 =
|
|
SetKernelDistOutput(&std::get<1>(api_output), spmd_info.second[1]);
|
|
}
|
|
auto dense_out_1 =
|
|
dist_out_1 ? dist_out_1->unsafe_mutable_value() : nullptr;
|
|
if (!rank_is_in_current_mesh) {
|
|
*dense_out_1 =
|
|
phi::DenseTensor(std::make_shared<phi::Allocation>(
|
|
nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}
|
|
|
|
// 3. Infer DistTensor's Global Shape
|
|
phi::MetaTensor meta_dist_out_0(dist_out_0);
|
|
phi::MetaTensor meta_dist_out_1(dist_out_1);
|
|
phi::FusedGemmEpilogueInferMeta(MakeMetaTensor(*x.impl()),
|
|
MakeMetaTensor(*y.impl()),
|
|
MakeMetaTensor(*bias.impl()),
|
|
trans_x,
|
|
trans_y,
|
|
activation,
|
|
dist_out_0 ? &meta_dist_out_0 : nullptr,
|
|
dist_out_1 ? &meta_dist_out_1 : nullptr);
|
|
|
|
if (rank_is_in_current_mesh) {
|
|
// 4. Select Kernel
|
|
VLOG(6) << "fused_gemm_epilogue API dist branch: kernel key: ["
|
|
<< kernel_backend << ", " << kernel_layout << ", "
|
|
<< kernel_data_type << "]";
|
|
auto kernel_result =
|
|
phi::KernelFactory::Instance().SelectKernelOrThrowError(
|
|
"fused_gemm_epilogue",
|
|
{kernel_backend, kernel_layout, kernel_data_type});
|
|
const auto& kernel = kernel_result.kernel;
|
|
VLOG(6) << "fused_gemm_epilogue kernel: " << kernel;
|
|
dev_ctx = GetDeviceContextByBackend(
|
|
kernel_result.has_fallback_cpu ? Backend::CPU : kernel_backend);
|
|
|
|
// 5. Reshard Input
|
|
auto dist_input_x =
|
|
ReshardApiInputToKernelInput(dev_ctx, x, spmd_info.first[0], "x");
|
|
auto dist_input_y =
|
|
ReshardApiInputToKernelInput(dev_ctx, y, spmd_info.first[1], "y");
|
|
auto dist_input_bias = ReshardApiInputToKernelInput(
|
|
dev_ctx, bias, spmd_info.first[2], "bias");
|
|
|
|
// 6. PrepareData (DataTransform & Prepare Dense Input)
|
|
dist_input_x = PrepareDataForDistTensor(
|
|
dist_input_x,
|
|
GetKernelInputArgDef(kernel.InputAt(0), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_x = &dist_input_x->value();
|
|
|
|
dist_input_y = PrepareDataForDistTensor(
|
|
dist_input_y,
|
|
GetKernelInputArgDef(kernel.InputAt(1), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_y = &dist_input_y->value();
|
|
|
|
dist_input_bias = PrepareDataForDistTensor(
|
|
dist_input_bias,
|
|
GetKernelInputArgDef(kernel.InputAt(2), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_bias = &dist_input_bias->value();
|
|
|
|
// 7. RecordOpInfoSupplement
|
|
if (phi::RecordOpInfoSupplement::IsEnabled()) {
|
|
std::vector<std::pair<const char*, std::vector<phi::DDim>>>
|
|
input_shapes{{"x", {(*input_x).dims()}},
|
|
{"y", {(*input_y).dims()}},
|
|
{"bias", {(*input_bias).dims()}}};
|
|
phi::AttributeMap attrs;
|
|
attrs["trans_x"] = trans_x;
|
|
attrs["trans_y"] = trans_y;
|
|
attrs["activation"] = activation;
|
|
phi::RecordOpInfoSupplement("fused_gemm_epilogue", input_shapes, attrs);
|
|
}
|
|
// 8. Infer Local DenseTensor Meta
|
|
phi::MetaTensor meta_dense_out_0(dense_out_0);
|
|
phi::MetaTensor meta_dense_out_1(dense_out_1);
|
|
phi::FusedGemmEpilogueInferMeta(
|
|
MakeMetaTensor(*input_x),
|
|
MakeMetaTensor(*input_y),
|
|
MakeMetaTensor(*input_bias),
|
|
trans_x,
|
|
trans_y,
|
|
activation,
|
|
dense_out_0 ? &meta_dense_out_0 : nullptr,
|
|
dense_out_1 ? &meta_dense_out_1 : nullptr);
|
|
|
|
// 9. DenseTensor Kernel Call
|
|
phi::RecordEvent* kernel_record_event = nullptr;
|
|
if (phi::RecordEvent::IsEnabled()) {
|
|
kernel_record_event =
|
|
new phi::RecordEvent("fused_gemm_epilogue dist compute",
|
|
phi::TracerEventType::DygraphKernelLaunch,
|
|
1);
|
|
}
|
|
using kernel_signature = void (*)(const phi::DeviceContext&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
bool,
|
|
bool,
|
|
const std::string&,
|
|
phi::DenseTensor*,
|
|
phi::DenseTensor*);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
(*kernel_fn)(*dev_ctx,
|
|
*input_x,
|
|
*input_y,
|
|
*input_bias,
|
|
trans_x,
|
|
trans_y,
|
|
activation,
|
|
dense_out_0,
|
|
dense_out_1);
|
|
if (FLAGS_benchmark) {
|
|
dev_ctx->Wait();
|
|
std::cout << "fused_gemm_epilogue kernel run finish." << std::endl;
|
|
}
|
|
if (kernel_record_event != nullptr) {
|
|
delete kernel_record_event;
|
|
}
|
|
|
|
// 10. Fallback
|
|
if (kernel_result.has_fallback_cpu) {
|
|
TransDataBackend(dense_out_0, kernel_backend, dense_out_0);
|
|
TransDataBackend(dense_out_1, kernel_backend, dense_out_1);
|
|
}
|
|
}
|
|
|
|
// 11. Set Output Dist Attr For Default Impl
|
|
// API `fused_gemm_epilogue` does not need to set DistAttr for output.
|
|
|
|
// 12. Return
|
|
return api_output;
|
|
}
|
|
#endif // PADDLE_WITH_DISTRIBUTE
|
|
|
|
VLOG(6) << "fused_gemm_epilogue API kernel key: [" << kernel_backend << ", "
|
|
<< kernel_layout << ", " << kernel_data_type << "]";
|
|
auto kernel_result = phi::KernelFactory::Instance().SelectKernelOrThrowError(
|
|
"fused_gemm_epilogue",
|
|
{kernel_backend, kernel_layout, kernel_data_type},
|
|
true);
|
|
const auto& kernel = kernel_result.kernel;
|
|
if (FLAGS_low_precision_op_list) {
|
|
phi::KernelFactory::Instance().AddToLowPrecisionKernelList(
|
|
"fused_gemm_epilogue", kernel_data_type);
|
|
}
|
|
VLOG(6) << "fused_gemm_epilogue kernel: " << kernel;
|
|
// add actual_kernel_backend to select actual kernel backend after a potential
|
|
// falling-back to CPU
|
|
Backend actual_kernel_backend =
|
|
kernel_result.has_fallback_cpu ? Backend::CPU : kernel_backend;
|
|
auto* dev_ctx = GetDeviceContextByBackend(actual_kernel_backend);
|
|
|
|
auto input_x = PrepareData(
|
|
x,
|
|
GetKernelInputArgDef(kernel.InputAt(0), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_y = PrepareData(
|
|
y,
|
|
GetKernelInputArgDef(kernel.InputAt(1), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_bias = PrepareData(
|
|
bias,
|
|
GetKernelInputArgDef(kernel.InputAt(2), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
if (phi::RecordOpInfoSupplement::IsEnabled()) {
|
|
std::vector<std::pair<const char*, std::vector<phi::DDim>>> input_shapes{
|
|
{"x", {(*input_x).dims()}},
|
|
{"y", {(*input_y).dims()}},
|
|
{"bias", {(*input_bias).dims()}}};
|
|
phi::AttributeMap attrs;
|
|
attrs["trans_x"] = trans_x;
|
|
attrs["trans_y"] = trans_y;
|
|
attrs["activation"] = activation;
|
|
phi::RecordOpInfoSupplement("fused_gemm_epilogue", input_shapes, attrs);
|
|
}
|
|
|
|
std::tuple<Tensor, Tensor> api_output;
|
|
auto kernel_out_0 = SetKernelOutput(&std::get<0>(api_output));
|
|
phi::DenseTensor* kernel_out_1 = nullptr;
|
|
if (activation != "none") {
|
|
kernel_out_1 = SetKernelOutput(&std::get<1>(api_output));
|
|
}
|
|
|
|
phi::RecordEvent* infer_shape_record_event = nullptr;
|
|
if (phi::RecordEvent::IsEnabled()) {
|
|
infer_shape_record_event =
|
|
new phi::RecordEvent("fused_gemm_epilogue infer_meta",
|
|
phi::TracerEventType::OperatorInner,
|
|
1);
|
|
}
|
|
phi::MetaTensor meta_out_0(kernel_out_0, kernel_result.is_stride_kernel);
|
|
phi::MetaTensor meta_out_1(kernel_out_1, kernel_result.is_stride_kernel);
|
|
|
|
std::vector<phi::MetaTensor*> output_metas_for_compact;
|
|
if (kernel_out_0) output_metas_for_compact.push_back(&meta_out_0);
|
|
if (kernel_out_1) output_metas_for_compact.push_back(&meta_out_1);
|
|
|
|
phi::FusedGemmEpilogueInferMeta(MakeMetaTensor(*input_x),
|
|
MakeMetaTensor(*input_y),
|
|
MakeMetaTensor(*input_bias),
|
|
trans_x,
|
|
trans_y,
|
|
activation,
|
|
kernel_out_0 ? &meta_out_0 : nullptr,
|
|
kernel_out_1 ? &meta_out_1 : nullptr);
|
|
CheckAndDoCompact(output_metas_for_compact, "fused_gemm_epilogue");
|
|
|
|
if (infer_shape_record_event != nullptr) {
|
|
delete infer_shape_record_event;
|
|
}
|
|
using kernel_signature = void (*)(const phi::DeviceContext&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
bool,
|
|
bool,
|
|
const std::string&,
|
|
phi::DenseTensor*,
|
|
phi::DenseTensor*);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
phi::RecordEvent* kernel_record_event = nullptr;
|
|
if (phi::RecordEvent::IsEnabled()) {
|
|
kernel_record_event =
|
|
new phi::RecordEvent("fused_gemm_epilogue kernel launch",
|
|
phi::TracerEventType::DygraphKernelLaunch,
|
|
1);
|
|
}
|
|
(*kernel_fn)(*dev_ctx,
|
|
*input_x,
|
|
*input_y,
|
|
*input_bias,
|
|
trans_x,
|
|
trans_y,
|
|
activation,
|
|
kernel_out_0,
|
|
kernel_out_1);
|
|
if (FLAGS_benchmark) {
|
|
dev_ctx->Wait();
|
|
std::cout << "fused_gemm_epilogue kernel run finish." << std::endl;
|
|
}
|
|
if (kernel_record_event != nullptr) {
|
|
delete kernel_record_event;
|
|
}
|
|
if (kernel_result.has_fallback_cpu) {
|
|
TransDataBackend(kernel_out_0, kernel_backend, kernel_out_0);
|
|
TransDataBackend(kernel_out_1, kernel_backend, kernel_out_1);
|
|
}
|
|
|
|
return api_output;
|
|
}
|
|
|
|
// weight_list.size() should be weight_list.get_ptr()->size() but can't modify
|
|
// yaml file
|
|
std::tuple<Tensor, Tensor, Tensor, std::vector<Tensor>> cudnn_lstm_grad_impl(
|
|
const Tensor& x,
|
|
const Tensor& init_h,
|
|
const Tensor& init_c,
|
|
const paddle::optional<std::vector<Tensor>>& weight_list,
|
|
const paddle::optional<Tensor>& sequence_length,
|
|
const Tensor& out,
|
|
const Tensor& reserve,
|
|
const Tensor& state_out,
|
|
const Tensor& out_grad,
|
|
const Tensor& last_h_grad,
|
|
const Tensor& last_c_grad,
|
|
float dropout_prob,
|
|
bool is_bidirec,
|
|
int hidden_size,
|
|
int num_layers,
|
|
bool is_test,
|
|
int seed) {
|
|
// Kernel Key Construction
|
|
Backend kernel_backend = Backend::UNDEFINED;
|
|
DataLayout kernel_layout = DataLayout::UNDEFINED;
|
|
DataType kernel_data_type = DataType::UNDEFINED;
|
|
|
|
#ifdef PADDLE_WITH_DISTRIBUTE
|
|
bool run_auto_parallel = AllInputsAreDistTensor(x,
|
|
init_h,
|
|
init_c,
|
|
weight_list,
|
|
sequence_length,
|
|
out,
|
|
reserve,
|
|
state_out,
|
|
out_grad,
|
|
last_h_grad,
|
|
last_c_grad);
|
|
bool rank_is_in_current_mesh = true;
|
|
if (run_auto_parallel) {
|
|
auto mesh = std::static_pointer_cast<phi::distributed::DistTensor>(
|
|
last_c_grad.impl())
|
|
->dist_attr()
|
|
.process_mesh();
|
|
rank_is_in_current_mesh = phi::distributed::IsCurRankInMesh(mesh);
|
|
}
|
|
if (rank_is_in_current_mesh) {
|
|
kernel_data_type = ParseDataType(out_grad);
|
|
|
|
if (kernel_backend == Backend::UNDEFINED ||
|
|
kernel_layout == DataLayout::UNDEFINED ||
|
|
kernel_data_type == DataType::UNDEFINED) {
|
|
auto kernel_key_set = ParseKernelKeyByInputArgs(x,
|
|
init_h,
|
|
init_c,
|
|
weight_list,
|
|
sequence_length,
|
|
out,
|
|
reserve,
|
|
state_out,
|
|
out_grad,
|
|
last_h_grad,
|
|
last_c_grad);
|
|
auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey();
|
|
if (kernel_backend == Backend::UNDEFINED) {
|
|
kernel_backend = kernel_key.backend();
|
|
}
|
|
if (kernel_layout == DataLayout::UNDEFINED) {
|
|
kernel_layout = kernel_key.layout();
|
|
}
|
|
if (kernel_data_type == DataType::UNDEFINED) {
|
|
kernel_data_type = kernel_key.dtype();
|
|
}
|
|
}
|
|
}
|
|
|
|
// Kernel Dispatch Body
|
|
// Auto Parallel condition
|
|
if (run_auto_parallel) {
|
|
// 1. InferSpmd (Infer DistAttr of Inputs&Outputs)
|
|
auto meta_dist_input_x = MakeDistMetaTensor(*x.impl());
|
|
auto meta_dist_input_init_h = MakeDistMetaTensor(*init_h.impl());
|
|
auto meta_dist_input_init_c = MakeDistMetaTensor(*init_c.impl());
|
|
std::vector<phi::distributed::DistMetaTensor> meta_dist_input_weight_list;
|
|
if (weight_list) {
|
|
for (auto& e : *weight_list) {
|
|
meta_dist_input_weight_list.push_back(MakeDistMetaTensor(*e.impl()));
|
|
}
|
|
}
|
|
auto meta_dist_input_sequence_length =
|
|
sequence_length ? MakeDistMetaTensor(*(*sequence_length).impl())
|
|
: phi::distributed::DistMetaTensor();
|
|
auto meta_dist_input_out = MakeDistMetaTensor(*out.impl());
|
|
auto meta_dist_input_reserve = MakeDistMetaTensor(*reserve.impl());
|
|
auto meta_dist_input_state_out = MakeDistMetaTensor(*state_out.impl());
|
|
auto meta_dist_input_out_grad = MakeDistMetaTensor(*out_grad.impl());
|
|
auto meta_dist_input_last_h_grad = MakeDistMetaTensor(*last_h_grad.impl());
|
|
auto meta_dist_input_last_c_grad = MakeDistMetaTensor(*last_c_grad.impl());
|
|
auto spmd_info = phi::distributed::VariadicReplicatedInferSpmdDynamic(
|
|
meta_dist_input_x,
|
|
meta_dist_input_init_h,
|
|
meta_dist_input_init_c,
|
|
meta_dist_input_weight_list,
|
|
meta_dist_input_sequence_length,
|
|
meta_dist_input_out,
|
|
meta_dist_input_reserve,
|
|
meta_dist_input_state_out,
|
|
meta_dist_input_out_grad,
|
|
meta_dist_input_last_h_grad,
|
|
meta_dist_input_last_c_grad);
|
|
DebugInfoForInferSpmd("cudnn_lstm_grad", spmd_info);
|
|
|
|
// 2. Create API Output & Prepare Dist and Dense Output
|
|
phi::DeviceContext* dev_ctx = nullptr;
|
|
std::tuple<Tensor, Tensor, Tensor, std::vector<Tensor>> api_output;
|
|
|
|
auto dist_out_0 = SetKernelDistOutput(&std::get<0>(api_output));
|
|
auto dense_out_0 =
|
|
dist_out_0 ? dist_out_0->unsafe_mutable_value() : nullptr;
|
|
if (!rank_is_in_current_mesh) {
|
|
*dense_out_0 =
|
|
phi::DenseTensor(std::make_shared<phi::Allocation>(
|
|
nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}
|
|
|
|
auto dist_out_1 = SetKernelDistOutput(&std::get<1>(api_output));
|
|
auto dense_out_1 =
|
|
dist_out_1 ? dist_out_1->unsafe_mutable_value() : nullptr;
|
|
if (!rank_is_in_current_mesh) {
|
|
*dense_out_1 =
|
|
phi::DenseTensor(std::make_shared<phi::Allocation>(
|
|
nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}
|
|
|
|
auto dist_out_2 = SetKernelDistOutput(&std::get<2>(api_output));
|
|
auto dense_out_2 =
|
|
dist_out_2 ? dist_out_2->unsafe_mutable_value() : nullptr;
|
|
if (!rank_is_in_current_mesh) {
|
|
*dense_out_2 =
|
|
phi::DenseTensor(std::make_shared<phi::Allocation>(
|
|
nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}
|
|
|
|
auto dist_out_3 = SetKernelDistOutput(weight_list.get_ptr()->size(),
|
|
&std::get<3>(api_output));
|
|
std::vector<phi::DenseTensor*> dense_out_3(dist_out_3.size());
|
|
for (size_t i = 0; i < dist_out_3.size(); ++i) {
|
|
dense_out_3[i] = const_cast<phi::DenseTensor*>(&dist_out_3[i]->value());
|
|
if (!rank_is_in_current_mesh) {
|
|
*dense_out_3[i] = phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(
|
|
nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}
|
|
}
|
|
|
|
// 3. Infer DistTensor's Global Shape
|
|
phi::MetaTensor meta_dist_out_0(dist_out_0);
|
|
phi::MetaTensor meta_dist_out_1(dist_out_1);
|
|
phi::MetaTensor meta_dist_out_2(dist_out_2);
|
|
std::vector<phi::MetaTensor> dist_out_3_meta_vec;
|
|
for (phi::distributed::DistTensor* tmp : dist_out_3) {
|
|
dist_out_3_meta_vec.emplace_back(phi::MetaTensor(tmp));
|
|
}
|
|
std::vector<phi::MetaTensor*> dist_out_3_meta_ptr_vec(dist_out_3.size());
|
|
for (size_t i = 0; i < dist_out_3_meta_vec.size(); ++i) {
|
|
dist_out_3_meta_ptr_vec[i] =
|
|
dist_out_3[i] ? &dist_out_3_meta_vec[i] : nullptr;
|
|
}
|
|
|
|
std::vector<phi::MetaTensor> weight_list_meta_vec_tmp;
|
|
if (weight_list) {
|
|
for (auto tmp : *weight_list) {
|
|
weight_list_meta_vec_tmp.emplace_back(MakeMetaTensor(*tmp.impl()));
|
|
}
|
|
}
|
|
std::vector<const phi::MetaTensor*> weight_list_meta_ptr_vec_tmp(
|
|
weight_list_meta_vec_tmp.size());
|
|
for (size_t i = 0; i < weight_list_meta_ptr_vec_tmp.size(); ++i) {
|
|
weight_list_meta_ptr_vec_tmp[i] = &weight_list_meta_vec_tmp[i];
|
|
}
|
|
paddle::optional<std::vector<const phi::MetaTensor*>>
|
|
weight_list_meta_ptr_vec =
|
|
weight_list
|
|
? paddle::make_optional<std::vector<const phi::MetaTensor*>>(
|
|
weight_list_meta_ptr_vec_tmp)
|
|
: paddle::none;
|
|
|
|
phi::CudnnLSTMGradInferMeta(MakeMetaTensor(*x.impl()),
|
|
MakeMetaTensor(*init_h.impl()),
|
|
MakeMetaTensor(*init_c.impl()),
|
|
weight_list_meta_ptr_vec,
|
|
dist_out_0 ? &meta_dist_out_0 : nullptr,
|
|
dist_out_1 ? &meta_dist_out_1 : nullptr,
|
|
dist_out_2 ? &meta_dist_out_2 : nullptr,
|
|
dist_out_3_meta_ptr_vec);
|
|
|
|
if (rank_is_in_current_mesh) {
|
|
// 4. Select Kernel
|
|
VLOG(6) << "cudnn_lstm_grad API dist branch: kernel key: ["
|
|
<< kernel_backend << ", " << kernel_layout << ", "
|
|
<< kernel_data_type << "]";
|
|
auto kernel_result =
|
|
phi::KernelFactory::Instance().SelectKernelOrThrowError(
|
|
"cudnn_lstm_grad",
|
|
{kernel_backend, kernel_layout, kernel_data_type});
|
|
const auto& kernel = kernel_result.kernel;
|
|
VLOG(6) << "cudnn_lstm_grad kernel: " << kernel;
|
|
dev_ctx = GetDeviceContextByBackend(
|
|
kernel_result.has_fallback_cpu ? Backend::CPU : kernel_backend);
|
|
|
|
// 5. Reshard Input
|
|
auto dist_input_x =
|
|
ReshardApiInputToKernelInput(dev_ctx, x, spmd_info.first[0], "x");
|
|
auto dist_input_init_h = ReshardApiInputToKernelInput(
|
|
dev_ctx, init_h, spmd_info.first[1], "init_h");
|
|
auto dist_input_init_c = ReshardApiInputToKernelInput(
|
|
dev_ctx, init_c, spmd_info.first[2], "init_c");
|
|
auto dist_input_weight_list = ReshardApiInputToKernelInput(
|
|
dev_ctx, weight_list, spmd_info.first[3], "weight_list");
|
|
auto dist_input_sequence_length = ReshardApiInputToKernelInput(
|
|
dev_ctx, sequence_length, spmd_info.first[4], "sequence_length");
|
|
auto dist_input_out =
|
|
ReshardApiInputToKernelInput(dev_ctx, out, spmd_info.first[5], "out");
|
|
auto dist_input_reserve = ReshardApiInputToKernelInput(
|
|
dev_ctx, reserve, spmd_info.first[6], "reserve");
|
|
auto dist_input_state_out = ReshardApiInputToKernelInput(
|
|
dev_ctx, state_out, spmd_info.first[7], "state_out");
|
|
auto dist_input_out_grad = ReshardApiInputToKernelInput(
|
|
dev_ctx, out_grad, spmd_info.first[8], "out_grad");
|
|
auto dist_input_last_h_grad = ReshardApiInputToKernelInput(
|
|
dev_ctx, last_h_grad, spmd_info.first[9], "last_h_grad");
|
|
auto dist_input_last_c_grad = ReshardApiInputToKernelInput(
|
|
dev_ctx, last_c_grad, spmd_info.first[10], "last_c_grad");
|
|
|
|
// 6. PrepareData (DataTransform & Prepare Dense Input)
|
|
dist_input_x = PrepareDataForDistTensor(
|
|
dist_input_x,
|
|
GetKernelInputArgDef(kernel.InputAt(0), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_x = &dist_input_x->value();
|
|
|
|
dist_input_init_h = PrepareDataForDistTensor(
|
|
dist_input_init_h,
|
|
GetKernelInputArgDef(kernel.InputAt(1), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_init_h = &dist_input_init_h->value();
|
|
|
|
dist_input_init_c = PrepareDataForDistTensor(
|
|
dist_input_init_c,
|
|
GetKernelInputArgDef(kernel.InputAt(2), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_init_c = &dist_input_init_c->value();
|
|
|
|
auto dist_input_weight_list_vec = PrepareDataForDistTensor(
|
|
dist_input_weight_list,
|
|
GetKernelInputArgDef(kernel.InputAt(3), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
std::vector<const phi::DenseTensor*> dense_input_weight_list_vec;
|
|
if (weight_list) {
|
|
for (auto tmp : *dist_input_weight_list_vec) {
|
|
dense_input_weight_list_vec.emplace_back(&tmp->value());
|
|
}
|
|
}
|
|
paddle::optional<std::vector<const phi::DenseTensor*>> input_weight_list(
|
|
dense_input_weight_list_vec);
|
|
std::vector<phi::MetaTensor> dense_input_weight_list_meta_vec =
|
|
MakeMetaTensor(dense_input_weight_list_vec);
|
|
std::vector<const phi::MetaTensor*>
|
|
dense_input_weight_list_meta_ptr_vec_tmp(
|
|
dense_input_weight_list_meta_vec.size());
|
|
for (size_t i = 0; i < dense_input_weight_list_meta_ptr_vec_tmp.size();
|
|
++i) {
|
|
dense_input_weight_list_meta_ptr_vec_tmp[i] =
|
|
&dense_input_weight_list_meta_vec[i];
|
|
}
|
|
paddle::optional<std::vector<const phi::MetaTensor*>>
|
|
dense_input_weight_list_meta_ptr_vec =
|
|
weight_list
|
|
? paddle::make_optional<std::vector<const phi::MetaTensor*>>(
|
|
dense_input_weight_list_meta_ptr_vec_tmp)
|
|
: paddle::none;
|
|
|
|
dist_input_sequence_length = PrepareDataForDistTensor(
|
|
dist_input_sequence_length,
|
|
GetKernelInputArgDef(kernel.InputAt(4), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
paddle::optional<phi::DenseTensor> input_sequence_length =
|
|
dist_input_sequence_length
|
|
? paddle::make_optional<phi::DenseTensor>(
|
|
(*dist_input_sequence_length)->value())
|
|
: paddle::none;
|
|
|
|
dist_input_out = PrepareDataForDistTensor(
|
|
dist_input_out,
|
|
GetKernelInputArgDef(kernel.InputAt(5), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_out = &dist_input_out->value();
|
|
|
|
dist_input_reserve = PrepareDataForDistTensor(
|
|
dist_input_reserve,
|
|
GetKernelInputArgDef(kernel.InputAt(6), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_reserve = &dist_input_reserve->value();
|
|
|
|
dist_input_state_out = PrepareDataForDistTensor(
|
|
dist_input_state_out,
|
|
GetKernelInputArgDef(kernel.InputAt(7), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_state_out = &dist_input_state_out->value();
|
|
|
|
dist_input_out_grad = PrepareDataForDistTensor(
|
|
dist_input_out_grad,
|
|
GetKernelInputArgDef(kernel.InputAt(8), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_out_grad = &dist_input_out_grad->value();
|
|
|
|
dist_input_last_h_grad = PrepareDataForDistTensor(
|
|
dist_input_last_h_grad,
|
|
GetKernelInputArgDef(kernel.InputAt(9), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_last_h_grad = &dist_input_last_h_grad->value();
|
|
|
|
dist_input_last_c_grad = PrepareDataForDistTensor(
|
|
dist_input_last_c_grad,
|
|
GetKernelInputArgDef(kernel.InputAt(10), kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_last_c_grad = &dist_input_last_c_grad->value();
|
|
|
|
// 7. RecordOpInfoSupplement
|
|
if (phi::RecordOpInfoSupplement::IsEnabled()) {
|
|
std::vector<phi::DDim> sequence_length_record_shapes;
|
|
if (input_sequence_length) {
|
|
sequence_length_record_shapes.push_back(
|
|
(*input_sequence_length).dims());
|
|
}
|
|
std::vector<std::pair<const char*, std::vector<phi::DDim>>>
|
|
input_shapes{{"x", {(*input_x).dims()}},
|
|
{"init_h", {(*input_init_h).dims()}},
|
|
{"init_c", {(*input_init_c).dims()}},
|
|
{"sequence_length", sequence_length_record_shapes},
|
|
{"out", {(*input_out).dims()}},
|
|
{"reserve", {(*input_reserve).dims()}},
|
|
{"state_out", {(*input_state_out).dims()}},
|
|
{"out_grad", {(*input_out_grad).dims()}},
|
|
{"last_h_grad", {(*input_last_h_grad).dims()}},
|
|
{"last_c_grad", {(*input_last_c_grad).dims()}}};
|
|
std::vector<phi::DDim> ddims_vec;
|
|
ddims_vec.clear();
|
|
if (input_weight_list) {
|
|
ddims_vec.reserve(input_weight_list->size());
|
|
for (size_t i = 0; i < input_weight_list->size(); ++i) {
|
|
ddims_vec.emplace_back((*input_weight_list->at(i)).dims());
|
|
}
|
|
}
|
|
input_shapes.emplace_back("weight_list", ddims_vec);
|
|
phi::AttributeMap attrs;
|
|
attrs["dropout_prob"] = dropout_prob;
|
|
attrs["is_bidirec"] = is_bidirec;
|
|
attrs["hidden_size"] = hidden_size;
|
|
attrs["num_layers"] = num_layers;
|
|
attrs["is_test"] = is_test;
|
|
attrs["seed"] = seed;
|
|
phi::RecordOpInfoSupplement("cudnn_lstm_grad", input_shapes, attrs);
|
|
}
|
|
// 8. Infer Local DenseTensor Meta
|
|
phi::MetaTensor meta_dense_out_0(dense_out_0);
|
|
phi::MetaTensor meta_dense_out_1(dense_out_1);
|
|
phi::MetaTensor meta_dense_out_2(dense_out_2);
|
|
std::vector<phi::MetaTensor> dense_out_3_meta_vec =
|
|
MakeMetaTensor(dense_out_3);
|
|
std::vector<phi::MetaTensor*> dense_out_3_meta_ptr_vec(
|
|
dense_out_3_meta_vec.size());
|
|
for (size_t i = 0; i < dense_out_3_meta_vec.size(); ++i) {
|
|
dense_out_3_meta_ptr_vec[i] =
|
|
dense_out_3[i] ? &dense_out_3_meta_vec[i] : nullptr;
|
|
}
|
|
|
|
phi::CudnnLSTMGradInferMeta(MakeMetaTensor(*input_x),
|
|
MakeMetaTensor(*input_init_h),
|
|
MakeMetaTensor(*input_init_c),
|
|
dense_input_weight_list_meta_ptr_vec,
|
|
dense_out_0 ? &meta_dense_out_0 : nullptr,
|
|
dense_out_1 ? &meta_dense_out_1 : nullptr,
|
|
dense_out_2 ? &meta_dense_out_2 : nullptr,
|
|
dense_out_3_meta_ptr_vec);
|
|
|
|
// 9. DenseTensor Kernel Call
|
|
phi::RecordEvent* kernel_record_event = nullptr;
|
|
if (phi::RecordEvent::IsEnabled()) {
|
|
kernel_record_event =
|
|
new phi::RecordEvent("cudnn_lstm_grad dist compute",
|
|
phi::TracerEventType::DygraphKernelLaunch,
|
|
1);
|
|
}
|
|
using kernel_signature = void (*)(
|
|
const phi::DeviceContext&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const paddle::optional<std::vector<const phi::DenseTensor*>>&,
|
|
const paddle::optional<phi::DenseTensor>&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
float,
|
|
bool,
|
|
int,
|
|
int,
|
|
bool,
|
|
int,
|
|
phi::DenseTensor*,
|
|
phi::DenseTensor*,
|
|
phi::DenseTensor*,
|
|
std::vector<phi::DenseTensor*>);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
(*kernel_fn)(*dev_ctx,
|
|
*input_x,
|
|
*input_init_h,
|
|
*input_init_c,
|
|
input_weight_list,
|
|
input_sequence_length,
|
|
*input_out,
|
|
*input_reserve,
|
|
*input_state_out,
|
|
*input_out_grad,
|
|
*input_last_h_grad,
|
|
*input_last_c_grad,
|
|
dropout_prob,
|
|
is_bidirec,
|
|
hidden_size,
|
|
num_layers,
|
|
is_test,
|
|
seed,
|
|
dense_out_0,
|
|
dense_out_1,
|
|
dense_out_2,
|
|
dense_out_3);
|
|
if (FLAGS_benchmark) {
|
|
dev_ctx->Wait();
|
|
std::cout << "cudnn_lstm_grad kernel run finish." << std::endl;
|
|
}
|
|
if (kernel_record_event != nullptr) {
|
|
delete kernel_record_event;
|
|
}
|
|
|
|
// 10. Fallback
|
|
if (kernel_result.has_fallback_cpu) {
|
|
TransDataBackend(dense_out_0, kernel_backend, dense_out_0);
|
|
TransDataBackend(dense_out_1, kernel_backend, dense_out_1);
|
|
TransDataBackend(dense_out_2, kernel_backend, dense_out_2);
|
|
TransDataBackend(dense_out_3, kernel_backend, dense_out_3);
|
|
}
|
|
}
|
|
|
|
// 11. Set Output Dist Attr For Default Impl
|
|
auto current_process_mesh =
|
|
paddle::holds_alternative<phi::distributed::TensorDistAttr>(
|
|
spmd_info.first[0])
|
|
? paddle::get<0>(spmd_info.first[0]).process_mesh()
|
|
: paddle::get<1>(spmd_info.first[0]).at(0).process_mesh();
|
|
SetReplicatedDistAttrForOutput(dist_out_0, current_process_mesh);
|
|
SetReplicatedDistAttrForOutput(dist_out_1, current_process_mesh);
|
|
SetReplicatedDistAttrForOutput(dist_out_2, current_process_mesh);
|
|
for (size_t i = 0; i < dist_out_3.size(); ++i) {
|
|
SetReplicatedDistAttrForOutput(dist_out_3[i], current_process_mesh);
|
|
}
|
|
|
|
// 12. Return
|
|
return api_output;
|
|
}
|
|
#endif // PADDLE_WITH_DISTRIBUTE
|
|
VLOG(6) << "cudnn_lstm_grad API kernel key: [" << kernel_backend << ", "
|
|
<< kernel_layout << ", " << kernel_data_type << "]";
|
|
auto kernel_result = phi::KernelFactory::Instance().SelectKernelOrThrowError(
|
|
"cudnn_lstm_grad",
|
|
{kernel_backend, kernel_layout, kernel_data_type},
|
|
true);
|
|
const auto& kernel = kernel_result.kernel;
|
|
if (FLAGS_low_precision_op_list) {
|
|
phi::KernelFactory::Instance().AddToLowPrecisionKernelList(
|
|
"cudnn_lstm_grad", kernel_data_type);
|
|
}
|
|
VLOG(6) << "cudnn_lstm_grad kernel: " << kernel;
|
|
// add actual_kernel_backend to select actual kernel backend after a potential
|
|
// falling-back to CPU
|
|
Backend actual_kernel_backend =
|
|
kernel_result.has_fallback_cpu ? Backend::CPU : kernel_backend;
|
|
auto* dev_ctx = GetDeviceContextByBackend(actual_kernel_backend);
|
|
|
|
auto input_x = PrepareData(
|
|
x,
|
|
GetKernelInputArgDef(kernel.InputAt(0), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_init_h = PrepareData(
|
|
init_h,
|
|
GetKernelInputArgDef(kernel.InputAt(1), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_init_c = PrepareData(
|
|
init_c,
|
|
GetKernelInputArgDef(kernel.InputAt(2), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_weight_list_vec = PrepareData(
|
|
weight_list,
|
|
GetKernelInputArgDef(kernel.InputAt(3), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
paddle::optional<std::vector<const phi::DenseTensor*>> input_weight_list;
|
|
if (input_weight_list_vec) {
|
|
input_weight_list = paddle::optional<std::vector<const phi::DenseTensor*>>(
|
|
input_weight_list_vec->size());
|
|
for (size_t i = 0; i < input_weight_list_vec->size(); ++i) {
|
|
input_weight_list->at(i) = &input_weight_list_vec->at(i);
|
|
}
|
|
}
|
|
auto input_sequence_length = PrepareData(
|
|
sequence_length,
|
|
GetKernelInputArgDef(kernel.InputAt(4), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_out = PrepareData(
|
|
out,
|
|
GetKernelInputArgDef(kernel.InputAt(5), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_reserve = PrepareData(
|
|
reserve,
|
|
GetKernelInputArgDef(kernel.InputAt(6), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_state_out = PrepareData(
|
|
state_out,
|
|
GetKernelInputArgDef(kernel.InputAt(7), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_out_grad = PrepareData(
|
|
out_grad,
|
|
GetKernelInputArgDef(kernel.InputAt(8), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_last_h_grad = PrepareData(
|
|
last_h_grad,
|
|
GetKernelInputArgDef(kernel.InputAt(9), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_last_c_grad = PrepareData(
|
|
last_c_grad,
|
|
GetKernelInputArgDef(kernel.InputAt(10), actual_kernel_backend),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
if (phi::RecordOpInfoSupplement::IsEnabled()) {
|
|
std::vector<phi::DDim> sequence_length_record_shapes;
|
|
if (input_sequence_length) {
|
|
sequence_length_record_shapes.push_back((*input_sequence_length).dims());
|
|
}
|
|
std::vector<std::pair<const char*, std::vector<phi::DDim>>> input_shapes{
|
|
{"x", {(*input_x).dims()}},
|
|
{"init_h", {(*input_init_h).dims()}},
|
|
{"init_c", {(*input_init_c).dims()}},
|
|
{"sequence_length", sequence_length_record_shapes},
|
|
{"out", {(*input_out).dims()}},
|
|
{"reserve", {(*input_reserve).dims()}},
|
|
{"state_out", {(*input_state_out).dims()}},
|
|
{"out_grad", {(*input_out_grad).dims()}},
|
|
{"last_h_grad", {(*input_last_h_grad).dims()}},
|
|
{"last_c_grad", {(*input_last_c_grad).dims()}}};
|
|
std::vector<phi::DDim> ddims_vec;
|
|
ddims_vec.clear();
|
|
if (input_weight_list) {
|
|
ddims_vec.reserve(input_weight_list->size());
|
|
for (size_t i = 0; i < input_weight_list->size(); ++i) {
|
|
ddims_vec.emplace_back((*input_weight_list->at(i)).dims());
|
|
}
|
|
}
|
|
input_shapes.emplace_back("weight_list", ddims_vec);
|
|
phi::AttributeMap attrs;
|
|
attrs["dropout_prob"] = dropout_prob;
|
|
attrs["is_bidirec"] = is_bidirec;
|
|
attrs["hidden_size"] = hidden_size;
|
|
attrs["num_layers"] = num_layers;
|
|
attrs["is_test"] = is_test;
|
|
attrs["seed"] = seed;
|
|
phi::RecordOpInfoSupplement("cudnn_lstm_grad", input_shapes, attrs);
|
|
}
|
|
|
|
std::tuple<Tensor, Tensor, Tensor, std::vector<Tensor>> api_output;
|
|
auto kernel_out_0 = SetKernelOutput(&std::get<0>(api_output));
|
|
auto kernel_out_1 = SetKernelOutput(&std::get<1>(api_output));
|
|
auto kernel_out_2 = SetKernelOutput(&std::get<2>(api_output));
|
|
auto kernel_out_3 =
|
|
SetKernelOutput(weight_list.get_ptr()->size(), &std::get<3>(api_output));
|
|
|
|
phi::RecordEvent* infer_shape_record_event = nullptr;
|
|
if (phi::RecordEvent::IsEnabled()) {
|
|
infer_shape_record_event = new phi::RecordEvent(
|
|
"cudnn_lstm_grad infer_meta", phi::TracerEventType::OperatorInner, 1);
|
|
}
|
|
|
|
auto weight_list_meta_vec = MakeMetaTensor(input_weight_list);
|
|
paddle::optional<std::vector<const phi::MetaTensor*>> weight_list_metas(
|
|
weight_list_meta_vec.size());
|
|
for (size_t i = 0; i < weight_list_meta_vec.size(); ++i) {
|
|
weight_list_metas->at(i) = &weight_list_meta_vec[i];
|
|
}
|
|
phi::MetaTensor meta_out_0(kernel_out_0, kernel_result.is_stride_kernel);
|
|
phi::MetaTensor meta_out_1(kernel_out_1, kernel_result.is_stride_kernel);
|
|
phi::MetaTensor meta_out_2(kernel_out_2, kernel_result.is_stride_kernel);
|
|
|
|
auto kernel_out_3_meta_vec = MakeMetaTensor(kernel_out_3);
|
|
std::vector<phi::MetaTensor*> kernel_out_3_metas(
|
|
kernel_out_3_meta_vec.size());
|
|
for (size_t i = 0; i < kernel_out_3_meta_vec.size(); ++i) {
|
|
kernel_out_3_metas[i] =
|
|
kernel_out_3[i] ? &kernel_out_3_meta_vec[i] : nullptr;
|
|
}
|
|
|
|
std::vector<phi::MetaTensor*> output_metas_for_compact;
|
|
if (kernel_out_0) output_metas_for_compact.push_back(&meta_out_0);
|
|
if (kernel_out_1) output_metas_for_compact.push_back(&meta_out_1);
|
|
if (kernel_out_2) output_metas_for_compact.push_back(&meta_out_1);
|
|
output_metas_for_compact.insert(output_metas_for_compact.end(),
|
|
kernel_out_3_metas.begin(),
|
|
kernel_out_3_metas.end());
|
|
|
|
phi::CudnnLSTMGradInferMeta(MakeMetaTensor(*input_x),
|
|
MakeMetaTensor(*input_init_h),
|
|
MakeMetaTensor(*input_init_c),
|
|
weight_list_metas,
|
|
kernel_out_0 ? &meta_out_0 : nullptr,
|
|
kernel_out_1 ? &meta_out_1 : nullptr,
|
|
kernel_out_2 ? &meta_out_2 : nullptr,
|
|
kernel_out_3_metas);
|
|
CheckAndDoCompact(output_metas_for_compact, "cudnn_lstm_grad");
|
|
|
|
if (infer_shape_record_event != nullptr) {
|
|
delete infer_shape_record_event;
|
|
}
|
|
using kernel_signature =
|
|
void (*)(const phi::DeviceContext&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const paddle::optional<std::vector<const phi::DenseTensor*>>&,
|
|
const paddle::optional<phi::DenseTensor>&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
float,
|
|
bool,
|
|
int,
|
|
int,
|
|
bool,
|
|
int,
|
|
phi::DenseTensor*,
|
|
phi::DenseTensor*,
|
|
phi::DenseTensor*,
|
|
std::vector<phi::DenseTensor*>);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
phi::RecordEvent* kernel_record_event = nullptr;
|
|
if (phi::RecordEvent::IsEnabled()) {
|
|
kernel_record_event =
|
|
new phi::RecordEvent("cudnn_lstm_grad kernel launch",
|
|
phi::TracerEventType::DygraphKernelLaunch,
|
|
1);
|
|
}
|
|
(*kernel_fn)(*dev_ctx,
|
|
*input_x,
|
|
*input_init_h,
|
|
*input_init_c,
|
|
input_weight_list,
|
|
input_sequence_length,
|
|
*input_out,
|
|
*input_reserve,
|
|
*input_state_out,
|
|
*input_out_grad,
|
|
*input_last_h_grad,
|
|
*input_last_c_grad,
|
|
dropout_prob,
|
|
is_bidirec,
|
|
hidden_size,
|
|
num_layers,
|
|
is_test,
|
|
seed,
|
|
kernel_out_0,
|
|
kernel_out_1,
|
|
kernel_out_2,
|
|
kernel_out_3);
|
|
if (FLAGS_benchmark) {
|
|
dev_ctx->Wait();
|
|
std::cout << "cudnn_lstm_grad kernel run finish." << std::endl;
|
|
}
|
|
if (kernel_record_event != nullptr) {
|
|
delete kernel_record_event;
|
|
}
|
|
if (kernel_result.has_fallback_cpu) {
|
|
TransDataBackend(kernel_out_0, kernel_backend, kernel_out_0);
|
|
TransDataBackend(kernel_out_1, kernel_backend, kernel_out_1);
|
|
TransDataBackend(kernel_out_2, kernel_backend, kernel_out_2);
|
|
TransDataBackend(kernel_out_3, kernel_backend, kernel_out_3);
|
|
}
|
|
|
|
return api_output;
|
|
}
|
|
|
|
////////////////// Backward(grad) api impls //////////////////////
|
|
|
|
void embedding_grad_impl(const Tensor& x,
|
|
const Tensor& weight,
|
|
const Tensor& out_grad,
|
|
int64_t padding_idx,
|
|
bool sparse,
|
|
Tensor* weight_grad) {
|
|
DataType kernel_data_type = ParseDataType(weight);
|
|
auto kernel_key_set = ParseKernelKeyByInputArgs(weight);
|
|
auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey();
|
|
VLOG(6) << "embedding_grad API kernel key: [" << kernel_key.backend() << ", "
|
|
<< kernel_key.layout() << ", " << kernel_data_type << "]";
|
|
|
|
if (phi::DenseTensor::classof(weight.impl().get()) ||
|
|
phi::distributed::DistTensor::classof(weight.impl().get())) {
|
|
std::string kernel_name =
|
|
sparse ? "embedding_sparse_grad" : "embedding_grad";
|
|
auto kernel_result =
|
|
phi::KernelFactory::Instance().SelectKernelOrThrowError(
|
|
kernel_name,
|
|
{kernel_key.backend(), kernel_key.layout(), kernel_data_type});
|
|
const auto& kernel = kernel_result.kernel;
|
|
VLOG(6) << kernel_name << " API kernel: " << kernel;
|
|
|
|
auto* dev_ctx = GetDeviceContextByBackend(
|
|
kernel_result.has_fallback_cpu ? Backend::CPU : kernel_key.backend());
|
|
|
|
#ifdef PADDLE_WITH_DISTRIBUTE
|
|
bool run_auto_parallel = AllInputsAreDistTensor(x, weight, out_grad);
|
|
// Auto Parallel condition
|
|
if (run_auto_parallel) {
|
|
bool rank_is_in_current_mesh = true;
|
|
auto mesh =
|
|
std::static_pointer_cast<phi::distributed::DistTensor>(x.impl())
|
|
->dist_attr()
|
|
.process_mesh();
|
|
rank_is_in_current_mesh = phi::distributed::IsCurRankInMesh(mesh);
|
|
|
|
// 1. InferSpmd (Infer DistAttr of Inputs&Outputs)
|
|
auto meta_dist_input_x = MakeDistMetaTensor(*x.impl());
|
|
auto meta_dist_input_weight = MakeDistMetaTensor(*weight.impl());
|
|
auto meta_dist_input_out_grad = MakeDistMetaTensor(*out_grad.impl());
|
|
auto spmd_info =
|
|
phi::distributed::EmbeddingGradInferSpmd(meta_dist_input_x,
|
|
meta_dist_input_weight,
|
|
meta_dist_input_out_grad,
|
|
padding_idx,
|
|
sparse);
|
|
|
|
// 2. Create Temporary Output & Prepare Dist and Dense Output
|
|
std::shared_ptr<phi::distributed::DistTensor> shared_dist_out =
|
|
CreateKernelDistOutput(
|
|
weight_grad, !rank_is_in_current_mesh, spmd_info.second[0]);
|
|
phi::distributed::DistTensor* dist_out = shared_dist_out.get();
|
|
phi::DenseTensor* dense_out = dist_out->unsafe_mutable_value();
|
|
if (dense_out && !rank_is_in_current_mesh && !dist_out->defined()) {
|
|
*dense_out = phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(
|
|
nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}
|
|
|
|
// 3. Infer DistTensor's Global Shape
|
|
phi::MetaTensor meta_dist_out(dist_out);
|
|
UnchangedInferMeta(MakeMetaTensor(*weight.impl()), &meta_dist_out);
|
|
|
|
// 4. Set Output Dist Attr For Default Impl
|
|
|
|
if (rank_is_in_current_mesh) {
|
|
// 5. Reshard Input
|
|
auto dist_input_x =
|
|
ReshardApiInputToKernelInput(dev_ctx, x, spmd_info.first[0]);
|
|
auto dist_input_weight =
|
|
ReshardApiInputToKernelInput(dev_ctx, weight, spmd_info.first[1]);
|
|
auto dist_input_out_grad =
|
|
ReshardApiInputToKernelInput(dev_ctx, out_grad, spmd_info.first[2]);
|
|
|
|
// 6. PrepareData (DataTransform & Prepare Dense Input)
|
|
dist_input_x = PrepareDataForDistTensor(
|
|
dist_input_x,
|
|
GetKernelInputArgDef(kernel.InputAt(0), kernel_key.backend()),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_x = &dist_input_x->value();
|
|
dist_input_weight = PrepareDataForDistTensor(
|
|
dist_input_weight,
|
|
GetKernelInputArgDef(kernel.InputAt(1), kernel_key.backend()),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_weight = &dist_input_weight->value();
|
|
dist_input_out_grad = PrepareDataForDistTensor(
|
|
dist_input_out_grad,
|
|
GetKernelInputArgDef(kernel.InputAt(2), kernel_key.backend()),
|
|
{},
|
|
kernel_result.is_stride_kernel);
|
|
auto input_out_grad = &dist_input_out_grad->value();
|
|
|
|
// 7. Infer Local DenseTensor Meta
|
|
phi::MetaTensor meta_dense_out(dense_out);
|
|
phi::EmbeddingGradInferMeta(MakeMetaTensor(*input_x),
|
|
MakeMetaTensor(*input_weight),
|
|
&meta_dense_out);
|
|
|
|
// 8. DenseTensor Kernel Call
|
|
using kernel_signature = void (*)(const phi::DeviceContext&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
int64_t,
|
|
phi::DenseTensor*);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
(*kernel_fn)(*dev_ctx,
|
|
*input_x,
|
|
*input_weight,
|
|
*input_out_grad,
|
|
padding_idx,
|
|
dense_out);
|
|
}
|
|
// 9. Reshard Kernel Output to API output
|
|
ReshardKernelOutputToApiOutput(dev_ctx, shared_dist_out, weight_grad);
|
|
|
|
// 10. Return
|
|
return;
|
|
}
|
|
#endif // PADDLE_WITH_DISTRIBUTE
|
|
|
|
auto input_x = PrepareData(x, kernel.InputAt(0), {}, false);
|
|
auto input_weight = PrepareData(weight, kernel.InputAt(1), {}, false);
|
|
auto input_out_grad = PrepareData(out_grad, kernel.InputAt(2), {}, false);
|
|
|
|
if (sparse) {
|
|
auto* kernel_out = SetSelectedRowsKernelOutput(weight_grad);
|
|
phi::MetaTensor meta_out(kernel_out);
|
|
meta_out.set_dims(input_weight->dims());
|
|
meta_out.set_dtype(input_weight->dtype());
|
|
kernel_out->set_height(input_weight->dims()[0]);
|
|
|
|
using kernel_signature = void (*)(const phi::DeviceContext&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
int64_t,
|
|
phi::SelectedRows*);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
(*kernel_fn)(*dev_ctx,
|
|
*input_x,
|
|
*input_weight,
|
|
*input_out_grad,
|
|
padding_idx,
|
|
kernel_out);
|
|
} else {
|
|
auto* kernel_out = SetKernelOutput(weight_grad);
|
|
phi::MetaTensor meta_out(kernel_out);
|
|
|
|
std::vector<phi::MetaTensor*> output_metas_for_compact;
|
|
output_metas_for_compact.push_back(&meta_out);
|
|
phi::UnchangedInferMeta(MakeMetaTensor(*input_weight), &meta_out);
|
|
CheckAndDoCompact(output_metas_for_compact, "embedding_grad");
|
|
using kernel_signature = void (*)(const phi::DeviceContext&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
const phi::DenseTensor&,
|
|
int64_t,
|
|
phi::DenseTensor*);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
(*kernel_fn)(*dev_ctx,
|
|
*input_x,
|
|
*input_weight,
|
|
*input_out_grad,
|
|
padding_idx,
|
|
kernel_out);
|
|
}
|
|
} else {
|
|
std::string kernel_name = sparse ? "sparse_weight_embedding_sparse_grad"
|
|
: "sparse_weight_embedding_grad";
|
|
auto kernel_result =
|
|
phi::KernelFactory::Instance().SelectKernelOrThrowError(
|
|
kernel_name,
|
|
{kernel_key.backend(), kernel_key.layout(), kernel_data_type});
|
|
const auto& kernel = kernel_result.kernel;
|
|
VLOG(6) << kernel_name << " API kernel: " << kernel;
|
|
|
|
auto* dev_ctx = GetDeviceContextByBackend(
|
|
kernel_result.has_fallback_cpu ? Backend::CPU : kernel_key.backend());
|
|
|
|
auto input_x = PrepareData(x, kernel.InputAt(0), {}, false);
|
|
auto input_weight = TensorToSelectedRows(weight);
|
|
auto input_out_grad = PrepareData(out_grad, kernel.InputAt(2), {}, false);
|
|
|
|
if (sparse) {
|
|
auto* kernel_out = SetSelectedRowsKernelOutput(weight_grad);
|
|
phi::MetaTensor meta_out(kernel_out);
|
|
phi::UnchangedInferMeta(MakeMetaTensor(*input_weight), &meta_out);
|
|
using kernel_signature = void (*)(const phi::DeviceContext&,
|
|
const phi::DenseTensor&,
|
|
const phi::SelectedRows&,
|
|
const phi::DenseTensor&,
|
|
int64_t,
|
|
phi::SelectedRows*);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
(*kernel_fn)(*dev_ctx,
|
|
*input_x,
|
|
*input_weight,
|
|
*input_out_grad,
|
|
padding_idx,
|
|
kernel_out);
|
|
} else {
|
|
auto* kernel_out = SetKernelOutput(weight_grad);
|
|
phi::MetaTensor meta_out(kernel_out);
|
|
meta_out.set_dims(input_weight->GetCompleteDims());
|
|
meta_out.set_dtype(input_weight->dtype());
|
|
using kernel_signature = void (*)(const phi::DeviceContext&,
|
|
const phi::DenseTensor&,
|
|
const phi::SelectedRows&,
|
|
const phi::DenseTensor&,
|
|
int64_t,
|
|
phi::DenseTensor*);
|
|
auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
(*kernel_fn)(*dev_ctx,
|
|
*input_x,
|
|
*input_weight,
|
|
*input_out_grad,
|
|
padding_idx,
|
|
kernel_out);
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace paddle::experimental
|