wehub-resource-sync
284 lines
11 KiB
C++
284 lines
11 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.
|
|
#pragma once
|
|
|
|
#include "glog/logging.h"
|
|
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/kernels/complex_kernel.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/impl/einsum_kernel_impl.h"
|
|
#include "paddle/phi/kernels/tile_grad_kernel.h"
|
|
#include "paddle/phi/kernels/tile_kernel.h"
|
|
#include "paddle/utils/string/string_helper.h"
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
DenseTensor PerformTileAndReduction(const Context& dev_ctx,
|
|
const LabelMap& label2type,
|
|
const LabelMap& label2shape,
|
|
const std::vector<int64_t>& broadcast_shape,
|
|
const std::vector<int64_t> x_shape,
|
|
std::string equ, // value pass
|
|
DenseTensor& t) { // NOLINT
|
|
auto tmp_label = equ;
|
|
auto tmp_union = unique_labels(tmp_label);
|
|
auto op_label = std::string(tmp_union.begin(), tmp_union.end());
|
|
VLOG(5) << "Start PerformTileAndReduction equation " << equ
|
|
<< " with operand shape: "
|
|
<< paddle::string::join_strings(vectorize<int64_t>(t.dims()), ",");
|
|
DenseTensor ret;
|
|
std::vector<int64_t> repeat_times;
|
|
std::vector<int64_t> resize_dims;
|
|
std::vector<int64_t> recover_shape;
|
|
std::vector<int64_t> t_shape = vectorize<int64_t>(t.dims());
|
|
for (size_t i = 0; i < op_label.size(); i++) {
|
|
int c = op_label[i];
|
|
if (label2type[c] == LabelType::Reduction) {
|
|
repeat_times.push_back(label2shape[c]);
|
|
resize_dims.push_back(1);
|
|
recover_shape.push_back(label2shape[c]);
|
|
t_shape.insert(t_shape.begin() + i, 1);
|
|
} else {
|
|
resize_dims.push_back(label2shape[c]);
|
|
repeat_times.push_back(1);
|
|
recover_shape.push_back(label2shape[c]);
|
|
}
|
|
}
|
|
PADDLE_ENFORCE_EQ(op_label.size(),
|
|
t_shape.size(),
|
|
common::errors::InvalidArgument(
|
|
"Input shape size doesn't match label nums, input "
|
|
"shape size: `%d`, but got label nums: `%d`",
|
|
t_shape.size(),
|
|
op_label.size()));
|
|
for (size_t i = 0; i < op_label.size(); i++) {
|
|
int c = op_label[i];
|
|
if (label2type[c] == LabelType::Contraction &&
|
|
t_shape[i] != label2shape[c]) {
|
|
repeat_times[i] = label2shape[c];
|
|
resize_dims[i] = 1;
|
|
}
|
|
}
|
|
t.Resize(resize_dims);
|
|
DenseTensor after_tile;
|
|
if (std::all_of(repeat_times.begin(), repeat_times.end(), [](int64_t x) {
|
|
return x == 1;
|
|
})) {
|
|
after_tile = t;
|
|
} else {
|
|
VLOG(4) << "do TileKernel with repeat_times="
|
|
<< paddle::string::join_strings(repeat_times, ",");
|
|
TileKernel<T, Context>(dev_ctx, t, repeat_times, &after_tile);
|
|
}
|
|
ret = after_tile;
|
|
VLOG(5) << "PermformTileAndReduction: recover shape: "
|
|
<< paddle::string::join_strings(recover_shape, ",");
|
|
ret.Resize(recover_shape);
|
|
|
|
// undiagonalize by einsum equation. only contain undiagonal operations.
|
|
DenseTensor undiagonal_out;
|
|
if (op_label != equ) {
|
|
VLOG(5) << "Undiagonal by einsum with args: " << op_label + "->" + equ;
|
|
EinsumInferKernel<T, Context>(
|
|
dev_ctx, {&ret}, op_label + "->" + equ, &undiagonal_out);
|
|
} else {
|
|
undiagonal_out = ret;
|
|
}
|
|
|
|
// call TileGradKernel to reverse broadcast operation.
|
|
VLOG(5) << "After diagonalize, we have tensor with shape: "
|
|
<< paddle::string::join_strings(
|
|
vectorize<int64_t>(undiagonal_out.dims()), ',');
|
|
repeat_times.clear();
|
|
for (size_t i = 0; i < x_shape.size(); ++i) {
|
|
VLOG(4) << "broadcast shape is " << broadcast_shape[i] << ", x_shape is "
|
|
<< x_shape[i];
|
|
repeat_times.push_back(broadcast_shape[i] / x_shape[i]);
|
|
}
|
|
bool is_all_ones = std::all_of(repeat_times.begin(),
|
|
repeat_times.end(),
|
|
[](int64_t x) { return x == 1; });
|
|
if (is_all_ones) {
|
|
VLOG(4) << "don't need broadcast recover, we just return undiagonal_out.";
|
|
return undiagonal_out;
|
|
}
|
|
DenseTensor tmp_x;
|
|
DenseTensor broadcast_out;
|
|
tmp_x.Resize(x_shape);
|
|
broadcast_out.Resize(x_shape);
|
|
TileGradKernel<T, Context>(
|
|
dev_ctx, tmp_x, undiagonal_out, repeat_times, &broadcast_out);
|
|
VLOG(5) << "After broadcast recover, we have tensor with shape: "
|
|
<< paddle::string::join_strings(
|
|
vectorize<int64_t>(broadcast_out.dims()), ',');
|
|
return broadcast_out;
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void EinsumGradKernel(const Context& dev_ctx,
|
|
const std::vector<const DenseTensor*>& x,
|
|
const std::vector<const DenseTensor*>& inner_cache,
|
|
const DenseTensor& out_grad,
|
|
const std::string& equation,
|
|
std::vector<DenseTensor*> x_grad) {
|
|
VLOG(5) << "Start EinsumGradKernel:";
|
|
bool has_zero_size_tensor = out_grad.numel() == 0;
|
|
for (auto& i : x_grad) {
|
|
if (i != nullptr) {
|
|
if (i->numel() == 0) {
|
|
has_zero_size_tensor = true;
|
|
}
|
|
Full<T, Context>(dev_ctx, i->dims(), 0, i);
|
|
}
|
|
}
|
|
if (has_zero_size_tensor) return;
|
|
LabelMap labelshape(0);
|
|
LabelMap labeltype(LabelType::Reduction);
|
|
std::vector<LabelMap> label2perms(x.size(), LabelMap(-1));
|
|
std::vector<char> all_labels; // order: ABO, AO, BO, AB, Reduce
|
|
std::vector<std::vector<int64_t>> broadcast_shapes(2);
|
|
std::vector<int64_t> output_dims;
|
|
|
|
std::vector<DDim> input_dims;
|
|
for (auto& i : x) {
|
|
input_dims.push_back(i->dims());
|
|
}
|
|
std::vector<std::string> input_strs;
|
|
std::string right;
|
|
ParseEinsumEquation(equation,
|
|
input_dims,
|
|
&labelshape,
|
|
&labeltype,
|
|
&all_labels,
|
|
&label2perms,
|
|
&broadcast_shapes,
|
|
&output_dims,
|
|
&right,
|
|
&input_strs);
|
|
|
|
VLOG(4) << "After grad parse einsum equation.";
|
|
|
|
auto gather_labels_except_reduction = [&labeltype](std::string all) {
|
|
std::string res("");
|
|
for (auto c : all)
|
|
if (labeltype[static_cast<int>(c)] != LabelType::Reduction) res += c;
|
|
auto tmp_unique = unique_labels(res);
|
|
return std::string(tmp_unique.begin(), tmp_unique.end());
|
|
};
|
|
if (x.size() == 1) { // Unary
|
|
auto splits = paddle::string::split_string(equation, "->");
|
|
auto left = splits[0];
|
|
right = splits[1];
|
|
auto new_equation = right + "->" + gather_labels_except_reduction(left);
|
|
auto new_operands = std::vector<const DenseTensor*>();
|
|
new_operands.push_back(&out_grad);
|
|
DenseTensor before_tile;
|
|
VLOG(5) << "new_equation is " << new_equation;
|
|
EinsumInferKernel<T, Context>(
|
|
dev_ctx, new_operands, new_equation, &before_tile);
|
|
*(x_grad[0]) =
|
|
PerformTileAndReduction<T, Context>(dev_ctx,
|
|
labeltype,
|
|
labelshape,
|
|
broadcast_shapes[0],
|
|
vectorize<int64_t>(x[0]->dims()),
|
|
left,
|
|
before_tile);
|
|
#ifndef PADDLE_WITH_XPU // xpu is not support conj now, we just disable it.
|
|
*(x_grad[0]) = Conj<T, Context>(dev_ctx, *x_grad[0]);
|
|
#endif
|
|
} else {
|
|
auto splits = paddle::string::split_string(equation, "->");
|
|
auto left = splits[0];
|
|
auto ops = paddle::string::split_string(left, ",");
|
|
right = splits[1];
|
|
auto equation_for_A =
|
|
ops[1] + "," + right + "->" + gather_labels_except_reduction(ops[0]);
|
|
auto equation_for_B =
|
|
right + "," + ops[0] + "->" + gather_labels_except_reduction(ops[1]);
|
|
auto operands_for_A = std::vector<const DenseTensor*>();
|
|
auto operands_for_B = std::vector<const DenseTensor*>();
|
|
DenseTensor dA, dB;
|
|
#ifndef PADDLE_WITH_XPU // xpu is not support conj now, we just disable it.
|
|
auto out_grad_conj = Conj<T, Context>(dev_ctx, out_grad);
|
|
#else
|
|
auto out_grad_conj = out_grad;
|
|
#endif
|
|
// dA = einsum(B, dC)
|
|
operands_for_A.push_back(x[1]);
|
|
operands_for_A.push_back(&out_grad_conj);
|
|
// dB = einsum(dC, A)
|
|
operands_for_B.push_back(&out_grad_conj);
|
|
operands_for_B.push_back(x[0]);
|
|
|
|
std::vector<DenseTensor> cache(3); // set empty; TA, TB, TdC
|
|
if (inner_cache.size() >
|
|
0) { // for compatibility, we can load and run v2.3 EinsumOp.
|
|
cache[0].ShareBufferWith(*(inner_cache[0]));
|
|
cache[1].ShareBufferWith(*(inner_cache[1]));
|
|
}
|
|
EinsumKernelImpl<T, Context>(dev_ctx,
|
|
all_labels,
|
|
labelshape,
|
|
operands_for_A,
|
|
equation_for_A,
|
|
&dA,
|
|
{&cache[1], &cache[2]},
|
|
false);
|
|
|
|
EinsumKernelImpl<T, Context>(dev_ctx,
|
|
all_labels,
|
|
labelshape,
|
|
operands_for_B,
|
|
equation_for_B,
|
|
&dB,
|
|
{&cache[2], &cache[0]},
|
|
false);
|
|
|
|
// release the cache tensor dTC to save memory right now. they are useless
|
|
// now.
|
|
cache.clear();
|
|
if (x_grad[0]) {
|
|
*(x_grad[0]) =
|
|
PerformTileAndReduction<T, Context>(dev_ctx,
|
|
labeltype,
|
|
labelshape,
|
|
broadcast_shapes[0],
|
|
vectorize<int64_t>(x[0]->dims()),
|
|
ops[0],
|
|
dA);
|
|
VLOG(4) << "After call dA";
|
|
#ifndef PADDLE_WITH_XPU // xpu is not support conj now, we just disable it.
|
|
*(x_grad[0]) = Conj<T, Context>(dev_ctx, *x_grad[0]);
|
|
#endif
|
|
}
|
|
if (x_grad[1]) {
|
|
*(x_grad[1]) =
|
|
PerformTileAndReduction<T, Context>(dev_ctx,
|
|
labeltype,
|
|
labelshape,
|
|
broadcast_shapes[1],
|
|
vectorize<int64_t>(x[1]->dims()),
|
|
ops[1],
|
|
dB);
|
|
#ifndef PADDLE_WITH_XPU // xpu is not support conj now, we just disable it.
|
|
*(x_grad[1]) = Conj<T, Context>(dev_ctx, *x_grad[1]);
|
|
#endif
|
|
VLOG(4) << "After call dA";
|
|
}
|
|
}
|
|
}
|
|
} // namespace phi
|