wehub-resource-sync
225 lines
8.2 KiB
C++
225 lines
8.2 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 "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/infermeta/binary.h"
|
|
#include "paddle/phi/infermeta/unary.h"
|
|
#include "paddle/phi/kernels/complex_kernel.h"
|
|
#include "paddle/phi/kernels/concat_kernel.h"
|
|
#include "paddle/phi/kernels/diagonal_kernel.h"
|
|
#include "paddle/phi/kernels/elementwise_add_kernel.h"
|
|
#include "paddle/phi/kernels/elementwise_subtract_kernel.h"
|
|
#include "paddle/phi/kernels/fill_diagonal_tensor_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/complex_functors.h"
|
|
#include "paddle/phi/kernels/funcs/for_range.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
#include "paddle/phi/kernels/funcs/parse_qr_mode.h"
|
|
#include "paddle/phi/kernels/matmul_kernel.h"
|
|
#include "paddle/phi/kernels/slice_kernel.h"
|
|
#include "paddle/phi/kernels/transpose_kernel.h"
|
|
#include "paddle/phi/kernels/triangular_solve_kernel.h"
|
|
#include "paddle/phi/kernels/tril_triu_kernel.h"
|
|
|
|
namespace phi {
|
|
|
|
template <class T, class Context>
|
|
static DenseTensor Fill(const Context& dev_ctx,
|
|
std::vector<int> shape,
|
|
float fill_value) {
|
|
DenseTensor ret;
|
|
ret.Resize(shape);
|
|
dev_ctx.template Alloc<T>(&ret);
|
|
funcs::SetConstant<Context, T>()(dev_ctx, &ret, T(fill_value));
|
|
return ret;
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void QrGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& q,
|
|
const DenseTensor& r,
|
|
const DenseTensor& q_grad,
|
|
const DenseTensor& r_grad,
|
|
const std::string& mode,
|
|
DenseTensor* x_grad) {
|
|
// Using alias names
|
|
const DenseTensor& A = x;
|
|
const DenseTensor& Q = q;
|
|
const DenseTensor& R = r;
|
|
const DenseTensor& dQ = q_grad;
|
|
const DenseTensor& dR = r_grad;
|
|
DenseTensor& dA = *x_grad;
|
|
|
|
dev_ctx.template Alloc<T>(&dA);
|
|
funcs::SetConstant<Context, T>()(dev_ctx, &dA, T(0));
|
|
|
|
bool compute_q, reduced;
|
|
std::tie(compute_q, reduced) = funcs::ParseQrMode(mode);
|
|
if (!compute_q) {
|
|
PADDLE_THROW(errors::InvalidArgument(
|
|
"The derivative of qr is not implemented when mode='%s'.", mode));
|
|
}
|
|
|
|
auto a_dims = A.dims();
|
|
int a_rank = a_dims.size();
|
|
int64_t m = a_dims[a_rank - 2];
|
|
int64_t n = a_dims[a_rank - 1];
|
|
|
|
if ((m > n) && (!reduced)) {
|
|
PADDLE_THROW(errors::InvalidArgument(
|
|
"The derivative of qr is not implemented when mode='complete' and "
|
|
"%d > %d.",
|
|
m,
|
|
n));
|
|
}
|
|
|
|
// m >= n case
|
|
auto m_ge_n_case = [](const Context& dev_ctx,
|
|
const DenseTensor& dQ,
|
|
const DenseTensor& dR,
|
|
const DenseTensor& A UNUSED,
|
|
const DenseTensor& Q,
|
|
const DenseTensor& R) -> DenseTensor {
|
|
// Roberts, D., & Roberts, L. (2020). QR and LQ Decomposition Matrix
|
|
// Backpropagation Algorithms for Square, Wide, and Deep Matrices and Their
|
|
// Software Implementation. https://arxiv.org/abs/2009.10071v4
|
|
|
|
// dR^H
|
|
DenseTensor R_term;
|
|
if (dR.initialized()) {
|
|
R_term = Matmul<T, Context>(dev_ctx,
|
|
R,
|
|
TransposeLast2Dim<T, Context>(
|
|
dev_ctx, Conj<T, Context>(dev_ctx, dR)));
|
|
} else {
|
|
R_term = Fill<T, Context>(dev_ctx, vectorize<int>(R.dims()), 0);
|
|
}
|
|
|
|
// dQ^H * Q
|
|
DenseTensor Q_term;
|
|
if (dQ.initialized()) {
|
|
Q_term = Matmul<T, Context>(
|
|
dev_ctx,
|
|
TransposeLast2Dim<T, Context>(dev_ctx, Conj<T, Context>(dev_ctx, dQ)),
|
|
Q);
|
|
} else {
|
|
Q_term = Fill<T, Context>(dev_ctx, vectorize<int>(R.dims()), 0);
|
|
}
|
|
|
|
DenseTensor M_tmp1 = Subtract<T, Context>(dev_ctx, R_term, Q_term);
|
|
DenseTensor M;
|
|
#ifdef PADDLE_WITH_HIP
|
|
// Compute M = (tril(M) + tril(M).mH()) * 0.5 Identity
|
|
DenseTensor M_tril_0 = TrilTriu<T, Context>(dev_ctx, M_tmp1, 0, true);
|
|
DenseTensor M_tril_1 = TrilTriu<T, Context>(dev_ctx, M_tmp1, -1, true);
|
|
M = Add<T, Context>(
|
|
dev_ctx, M_tril_0, TransposeLast2Dim<T, Context>(dev_ctx, M_tril_1));
|
|
#else
|
|
if (std::is_same<T, complex64>::value ||
|
|
std::is_same<T, complex128>::value) {
|
|
DenseTensor M_tril_tmp = TrilTriu<T, Context>(dev_ctx, M_tmp1, -1, true);
|
|
DenseTensor M_tril =
|
|
Add<T, Context>(dev_ctx,
|
|
M_tril_tmp,
|
|
TransposeLast2Dim<T, Context>(
|
|
dev_ctx, Conj<T, Context>(dev_ctx, M_tril_tmp)));
|
|
|
|
size_t rank = M_tmp1.dims().size();
|
|
DenseTensor M_diag_tmp =
|
|
Diagonal<T, Context>(dev_ctx, M_tmp1, 0, rank - 2, rank - 1);
|
|
DenseTensor M_diag_real = Real<T, Context>(dev_ctx, M_diag_tmp);
|
|
DenseTensor M_diag_imag = Fill<dtype::Real<T>, Context>(
|
|
dev_ctx, vectorize<int>(M_diag_real.dims()), 0);
|
|
|
|
DenseTensor M_diag;
|
|
M_diag.Resize(M_diag_real.dims());
|
|
dev_ctx.template Alloc<T>(&M_diag);
|
|
ComplexKernel<dtype::Real<T>>(dev_ctx, M_diag_real, M_diag_imag, &M_diag);
|
|
|
|
M = FillDiagonalTensor<T, Context>(
|
|
dev_ctx, M_tril, M_diag, 0, rank - 2, rank - 1);
|
|
} else {
|
|
// Compute M = (tril(M) + tril(M).mH()) * 0.5 Identity
|
|
DenseTensor M_tril_0 = TrilTriu<T, Context>(dev_ctx, M_tmp1, 0, true);
|
|
DenseTensor M_tril_1 = TrilTriu<T, Context>(dev_ctx, M_tmp1, -1, true);
|
|
M = Add<T, Context>(
|
|
dev_ctx, M_tril_0, TransposeLast2Dim<T, Context>(dev_ctx, M_tril_1));
|
|
}
|
|
#endif
|
|
|
|
DenseTensor rhs_term;
|
|
if (dQ.initialized()) {
|
|
rhs_term =
|
|
Add<T, Context>(dev_ctx, dQ, Matmul<T, Context>(dev_ctx, Q, M));
|
|
} else {
|
|
rhs_term = Matmul<T, Context>(dev_ctx, Q, M);
|
|
}
|
|
|
|
// dA * R^H = rhs_term
|
|
auto dA = TriangularSolve<T, Context>(
|
|
dev_ctx,
|
|
TransposeLast2Dim<T, Context>(
|
|
dev_ctx,
|
|
Conj<T, Context>(dev_ctx,
|
|
TransposeLast2Dim<T, Context>(dev_ctx, R))),
|
|
TransposeLast2Dim<T, Context>(dev_ctx, rhs_term),
|
|
/*upper=*/true,
|
|
/*transpose=*/false,
|
|
/*unitriangular=*/false);
|
|
|
|
return TransposeLast2Dim<T, Context>(dev_ctx, dA);
|
|
};
|
|
|
|
if (m >= n) {
|
|
auto dA_tmp = m_ge_n_case(dev_ctx, dQ, dR, A, Q, R);
|
|
Copy(dev_ctx, dA_tmp, dA.place(), false, &dA);
|
|
} else {
|
|
// If m < n for input matrices A, we partition A = [X|Y] and R = [U|V]
|
|
// Calculate dX and dY individually and concatenate them to get dA
|
|
dev_ctx.template Alloc<dtype::Real<T>>(&dA);
|
|
|
|
auto Y = Slice<T, Context>(dev_ctx, A, {A.dims().size() - 1}, {m}, {n});
|
|
auto U = Slice<T, Context>(dev_ctx, R, {R.dims().size() - 1}, {0}, {m});
|
|
DenseTensor dY, dX, dV, dU, dQ_prime;
|
|
|
|
if (dR.initialized()) {
|
|
dV = Slice<T, Context>(dev_ctx, dR, {dR.dims().size() - 1}, {m}, {n});
|
|
dU = Slice<T, Context>(dev_ctx, dR, {dR.dims().size() - 1}, {0}, {m});
|
|
// Y * dV^H
|
|
dQ_prime =
|
|
Matmul<T, Context>(dev_ctx,
|
|
Y,
|
|
TransposeLast2Dim<T, Context>(
|
|
dev_ctx, Conj<T, Context>(dev_ctx, dV)));
|
|
} else {
|
|
dV = Fill<T, Context>(dev_ctx, vectorize<int>(Y.dims()), 0);
|
|
dQ_prime = Fill<T, Context>(dev_ctx, vectorize<int>(Q.dims()), 0);
|
|
}
|
|
|
|
if (dQ.initialized()) {
|
|
dQ_prime = Add<T, Context>(dev_ctx, dQ, dQ_prime);
|
|
}
|
|
dX = m_ge_n_case(dev_ctx, dQ_prime, dU, A, Q, U);
|
|
dY = Matmul<T, Context>(dev_ctx, Q, dV);
|
|
// Concatenate dX and dY to get dA.
|
|
auto dA_tmp = Concat<T, Context>(dev_ctx, {&dX, &dY}, -1);
|
|
Copy(dev_ctx, dA_tmp, dA.place(), false, &dA);
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|