wehub-resource-sync
473 lines
16 KiB
C++
473 lines
16 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/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
|
namespace phi {
|
|
|
|
template <typename Context, typename T>
|
|
inline DenseTensor MatMul(const Context& dev_ctx,
|
|
const DenseTensor& matrix_a,
|
|
const DenseTensor& matrix_b,
|
|
const DDim& a_dim,
|
|
const DDim& b_dim) {
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
|
|
DenseTensor matrix_c;
|
|
DDim c_dim = make_ddim({a_dim[0], b_dim[1]});
|
|
matrix_c.Resize(c_dim);
|
|
dev_ctx.template Alloc<T>(&matrix_c);
|
|
|
|
auto mat_dim_a = funcs::CreateMatrixDescriptor(a_dim, 0, false);
|
|
auto mat_dim_b = funcs::CreateMatrixDescriptor(b_dim, 0, false);
|
|
const T alpha = static_cast<T>(1.0);
|
|
blas.MatMul(matrix_a.data<T>(),
|
|
mat_dim_a,
|
|
matrix_b.data<T>(),
|
|
mat_dim_b,
|
|
alpha,
|
|
matrix_c.data<T>(),
|
|
T(0));
|
|
return matrix_c;
|
|
}
|
|
|
|
/**
|
|
* @brief Recursively calculate matrix multiplication according to the optimal
|
|
* order
|
|
* Let k = order[i,j], then ins[i...j] = ins[i...k] * ins[k+1 ...j]
|
|
*
|
|
* @param
|
|
* ins: the input tensors
|
|
* ins_dims: the shape of ins after reshape
|
|
* order: the optimal order
|
|
* i: the left of sub chain
|
|
* j: the right of sub chain
|
|
* save_result: set true by backward
|
|
* results: save the intermediate result during backward
|
|
*/
|
|
template <typename Context, typename T>
|
|
inline DenseTensor MatChainMul(const Context& dev_ctx,
|
|
const std::vector<const DenseTensor*>& ins,
|
|
const std::vector<DDim>& ins_dims,
|
|
const std::vector<uint64_t>& order,
|
|
const uint64_t i,
|
|
const uint64_t j,
|
|
const bool save_result,
|
|
std::vector<DenseTensor>* results) {
|
|
if (i == j) {
|
|
return *ins[i];
|
|
}
|
|
|
|
const auto A = MatChainMul<Context, T>(dev_ctx,
|
|
ins,
|
|
ins_dims,
|
|
order,
|
|
i,
|
|
order[i * ins.size() + j],
|
|
save_result,
|
|
results);
|
|
DDim a_dim = A.dims();
|
|
if (i == order[i * ins.size() + j]) {
|
|
a_dim = ins_dims[i];
|
|
}
|
|
|
|
const auto B = MatChainMul<Context, T>(dev_ctx,
|
|
ins,
|
|
ins_dims,
|
|
order,
|
|
order[i * ins.size() + j] + 1,
|
|
j,
|
|
save_result,
|
|
results);
|
|
DDim b_dim = B.dims();
|
|
if (j == order[i * ins.size() + j] + 1) {
|
|
b_dim = ins_dims[j];
|
|
}
|
|
|
|
auto result = MatMul<Context, T>(dev_ctx, A, B, a_dim, b_dim);
|
|
if (save_result) {
|
|
(*results)[i * ins.size() + j] = result;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
/**
|
|
* @brief get the optimal order
|
|
*/
|
|
template <typename Context, typename T>
|
|
std::vector<uint64_t> GetOrder(const std::vector<const DenseTensor*>& ins,
|
|
const std::vector<DDim>& ins_dims) {
|
|
uint64_t n = ins.size();
|
|
// p: save the ins shape, the ins[i] shape is (p[i], p[i+1])
|
|
std::vector<uint64_t> p(n + 1);
|
|
for (uint64_t i = 0; i < n; i++) {
|
|
p[i] = ins_dims[i][0];
|
|
}
|
|
p[n] = ins_dims[n - 1][1];
|
|
|
|
// m[i, j]: save the lowest cost for multiplying ins[i...j]
|
|
std::vector<uint64_t> m(n * n, 0);
|
|
// define ins[i...j] means multiplying matrices from ins[i] to ins[j]
|
|
// order[i, j] = k, this means that ins[i...k] and ins[k...j] first and then
|
|
// multiply the resulting matrices is the optimal order for ins[i...j]
|
|
std::vector<uint64_t> order(n * n);
|
|
for (uint64_t l = 1; l < n; l++) {
|
|
for (uint64_t i = 0; i < n - l; i++) {
|
|
auto j = i + l;
|
|
m[i * n + j] = std::numeric_limits<uint64_t>::max();
|
|
for (uint64_t k = i; k < j; k++) {
|
|
uint64_t q =
|
|
m[i * n + k] + m[(k + 1) * n + j] + p[i] * p[k + 1] * p[j + 1];
|
|
if (q < m[i * n + j]) {
|
|
m[i * n + j] = q;
|
|
order[i * n + j] = k;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return order;
|
|
}
|
|
|
|
template <typename Context, typename T>
|
|
static inline DenseTensor MultiDotMatChainOrder(
|
|
const Context& dev_ctx,
|
|
const std::vector<const DenseTensor*>& ins,
|
|
const std::vector<DDim>& ins_dims,
|
|
const bool save_result,
|
|
std::vector<DenseTensor>* results) {
|
|
auto order = GetOrder<Context, T>(ins, ins_dims);
|
|
return MatChainMul<Context, T>(
|
|
dev_ctx, ins, ins_dims, order, 0, ins.size() - 1, save_result, results);
|
|
}
|
|
|
|
template <typename Context, typename T>
|
|
inline void GetDims(const std::vector<const DenseTensor*>& ins,
|
|
std::vector<DDim>* ins_dims) {
|
|
const auto n = ins.size();
|
|
for (size_t i = 0; i < n; i++) {
|
|
(*ins_dims)[i] = ins[i]->dims();
|
|
if (i == 0 && (*ins_dims)[i].size() == 1) {
|
|
(*ins_dims)[i] = make_ddim({1, (*ins_dims)[i][0]});
|
|
} else if (i == n - 1 && (*ins_dims)[i].size() == 1) {
|
|
(*ins_dims)[i] = make_ddim({(*ins_dims)[i][0], 1});
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void MultiDotKernel(const Context& dev_ctx,
|
|
const std::vector<const DenseTensor*>& x,
|
|
DenseTensor* out) {
|
|
auto ins = x;
|
|
dev_ctx.template Alloc<T>(out);
|
|
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
|
|
auto n = ins.size();
|
|
std::vector<DDim> ins_dims(n);
|
|
GetDims<Context, T>(ins, &ins_dims);
|
|
|
|
// If any numel is 0, then return.
|
|
bool size_0 = false;
|
|
for (size_t i = 0; i < n; i++) {
|
|
if (x[i]->numel() == 0) size_0 = true;
|
|
}
|
|
if (size_0) {
|
|
// For example: [2, 0], [0, 4] -> [2, 4]
|
|
if (out && out->numel() > 0) {
|
|
Full<T, Context>(dev_ctx, out->dims(), 0, out);
|
|
}
|
|
return;
|
|
}
|
|
const T scale = static_cast<T>(1.0);
|
|
if (n == 2) {
|
|
auto mat_dim_a = funcs::CreateMatrixDescriptor(ins_dims[0], 0, false);
|
|
auto mat_dim_b = funcs::CreateMatrixDescriptor(ins_dims[1], 0, false);
|
|
blas.MatMul(*ins[0], mat_dim_a, *ins[1], mat_dim_b, scale, out, T(0));
|
|
} else if (n == 3) {
|
|
const auto Ma = ins_dims[0][0];
|
|
const auto Ka = ins_dims[0][1];
|
|
const auto Nb = ins_dims[1][1];
|
|
const auto Nc = ins_dims[2][1];
|
|
const uint64_t cost1 = Ma * Nb * (Ka + Nc);
|
|
const uint64_t cost2 = Ka * Nc * (Nb + Ma);
|
|
auto mat_dim_a = funcs::CreateMatrixDescriptor(ins_dims[0], 0, false);
|
|
auto mat_dim_b = funcs::CreateMatrixDescriptor(ins_dims[1], 0, false);
|
|
auto mat_dim_c = funcs::CreateMatrixDescriptor(ins_dims[2], 0, false);
|
|
if (cost1 < cost2) {
|
|
DenseTensor tmp_out;
|
|
DDim tmp_dim = make_ddim({Ma, Nb});
|
|
tmp_out.Resize(tmp_dim);
|
|
dev_ctx.template Alloc<T>(&tmp_out);
|
|
blas.MatMul(
|
|
*ins[0], mat_dim_a, *ins[1], mat_dim_b, scale, &tmp_out, T(0));
|
|
auto mat_dim_tmp = funcs::CreateMatrixDescriptor(tmp_dim, 0, false);
|
|
blas.MatMul(tmp_out, mat_dim_tmp, *ins[2], mat_dim_c, scale, out, T(0));
|
|
} else {
|
|
DenseTensor tmp_out;
|
|
DDim tmp_dim = make_ddim({Ka, Nc});
|
|
tmp_out.Resize(tmp_dim);
|
|
dev_ctx.template Alloc<T>(&tmp_out);
|
|
blas.MatMul(
|
|
*ins[1], mat_dim_b, *ins[2], mat_dim_c, scale, &tmp_out, T(0));
|
|
auto mat_dim_tmp = funcs::CreateMatrixDescriptor(tmp_dim, 0, false);
|
|
blas.MatMul(*ins[0], mat_dim_a, tmp_out, mat_dim_tmp, scale, out, T(0));
|
|
}
|
|
} else {
|
|
std::vector<DenseTensor> results;
|
|
const auto tmp = MultiDotMatChainOrder<Context, T>(
|
|
dev_ctx, ins, ins_dims, false, &results);
|
|
auto out_dim = out->dims();
|
|
*out = tmp;
|
|
out->Resize(out_dim);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* @brief calculate dA and dB
|
|
* dA = dout * transpose(B)
|
|
* dB = transpose(A) * dout
|
|
*/
|
|
template <typename Context, typename T>
|
|
void CalcGrad(const Context& dev_ctx,
|
|
const DenseTensor& dout,
|
|
const DenseTensor& A,
|
|
const DenseTensor& B,
|
|
const DDim& dout_dim,
|
|
const DDim& a_dim,
|
|
const DDim& b_dim,
|
|
DenseTensor* dA,
|
|
DenseTensor* dB) {
|
|
auto mat_dim_dout = funcs::CreateMatrixDescriptor(dout_dim, 0, false);
|
|
auto mat_dim_a = funcs::CreateMatrixDescriptor(a_dim, 0, true);
|
|
auto mat_dim_b = funcs::CreateMatrixDescriptor(b_dim, 0, true);
|
|
T alpha = static_cast<T>(1.0);
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
blas.MatMul(A, mat_dim_a, dout, mat_dim_dout, alpha, dB, T(0));
|
|
blas.MatMul(dout, mat_dim_dout, B, mat_dim_b, alpha, dA, T(0));
|
|
}
|
|
|
|
/**
|
|
* @brief calculate multi matrix multiplication grad by a chain order
|
|
* @param
|
|
* dout: the grad of multi matrix multiplication out
|
|
* dx: the out grad of inputs
|
|
* ins: the input tensors
|
|
* ins_dims: the shape of ins after reshape
|
|
* order: the optimal order
|
|
* i: the left of sub chain
|
|
* j: the right of sub chain
|
|
* results: the intermediate result of forward
|
|
*/
|
|
template <typename Context, typename T>
|
|
void MatChainMulGrad(const Context& dev_ctx,
|
|
const DenseTensor& dout,
|
|
std::vector<DenseTensor*>* dx,
|
|
const std::vector<const DenseTensor*>& ins,
|
|
const DDim& dout_dim,
|
|
const std::vector<DDim>& ins_dims,
|
|
const std::vector<uint64_t>& order,
|
|
const uint64_t i,
|
|
const uint64_t j,
|
|
const std::vector<DenseTensor>& results) {
|
|
if (i == j) {
|
|
*((*dx)[i]) = dout;
|
|
return;
|
|
}
|
|
|
|
const auto n = ins.size();
|
|
const auto right = order[i * n + j];
|
|
const auto left = order[i * n + j] + 1;
|
|
// get the multi result of left sub chain
|
|
const auto* A = &results[i * n + right];
|
|
DDim a_dim = A->dims();
|
|
if (i == right) {
|
|
A = ins[i];
|
|
a_dim = ins_dims[i];
|
|
}
|
|
// get the multi result of right sub chain
|
|
const auto* B = &results[left * n + j];
|
|
DDim b_dim = B->dims();
|
|
if (left == j) {
|
|
B = ins[j];
|
|
b_dim = ins_dims[j];
|
|
}
|
|
DenseTensor dA, dB;
|
|
dA.Resize({dout_dim[0], b_dim[0]});
|
|
dB.Resize({a_dim[1], dout_dim[1]});
|
|
dev_ctx.template Alloc<T>(&dA);
|
|
dev_ctx.template Alloc<T>(&dB);
|
|
|
|
CalcGrad<Context, T>(dev_ctx, dout, *A, *B, dout_dim, a_dim, b_dim, &dA, &dB);
|
|
MatChainMulGrad<Context, T>(
|
|
dev_ctx, dA, dx, ins, dA.dims(), ins_dims, order, i, right, results);
|
|
MatChainMulGrad<Context, T>(
|
|
dev_ctx, dB, dx, ins, dB.dims(), ins_dims, order, left, j, results);
|
|
}
|
|
|
|
template <typename Context, typename T>
|
|
void MultiDotGradMatChainOrder(const Context& dev_ctx,
|
|
const DenseTensor& dout,
|
|
const std::vector<const DenseTensor*>& ins,
|
|
const DDim& dout_dim,
|
|
const std::vector<DDim>& ins_dims,
|
|
std::vector<DenseTensor*>* dx) {
|
|
auto order = GetOrder<Context, T>(ins, ins_dims);
|
|
auto n = ins.size();
|
|
std::vector<DenseTensor> results(static_cast<int64_t>(n) * n);
|
|
MatChainMul<Context, T>(
|
|
dev_ctx, ins, ins_dims, order, 0, n - 1, true, &results);
|
|
MatChainMulGrad<Context, T>(
|
|
dev_ctx, dout, dx, ins, dout_dim, ins_dims, order, 0, n - 1, results);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void MultiDotGradKernel(const Context& dev_ctx,
|
|
const std::vector<const DenseTensor*>& x,
|
|
const DenseTensor& out_grad,
|
|
std::vector<DenseTensor*> x_grad) {
|
|
auto ins = x;
|
|
auto dout = out_grad;
|
|
auto dx = x_grad;
|
|
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
|
|
bool size_0 = false;
|
|
const auto n = ins.size();
|
|
for (size_t i = 0; i < n; i++) {
|
|
dev_ctx.template Alloc<T>(dx[i]);
|
|
|
|
if (dx[i]->numel() == 0) {
|
|
size_0 = true;
|
|
}
|
|
}
|
|
if (size_0) {
|
|
for (size_t i = 0; i < n; i++) {
|
|
if (dx[i]->numel() > 0) {
|
|
Full<T, Context>(dev_ctx, dx[i]->dims(), 0, dx[i]);
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
std::vector<DDim> ins_dims(n);
|
|
GetDims<Context, T>(ins, &ins_dims);
|
|
|
|
DDim dout_dim = dout.dims();
|
|
if (ins[0]->dims().size() == 1 && ins[n - 1]->dims().size() == 1) {
|
|
dout_dim = make_ddim({1, 1});
|
|
} else if (ins[0]->dims().size() == 1) {
|
|
if (dout_dim.size() == 1) {
|
|
dout_dim = make_ddim({1, dout_dim[0]});
|
|
}
|
|
} else if (ins[n - 1]->dims().size() == 1) {
|
|
if (dout_dim.size() == 1) {
|
|
dout_dim = make_ddim({dout_dim[0], 1});
|
|
}
|
|
}
|
|
|
|
T alpha = static_cast<T>(1);
|
|
auto mat_dim_dout = funcs::CreateMatrixDescriptor(dout_dim, 0, false);
|
|
if (n == 2) {
|
|
CalcGrad<Context, T>(dev_ctx,
|
|
dout,
|
|
*ins[0],
|
|
*ins[1],
|
|
dout_dim,
|
|
ins_dims[0],
|
|
ins_dims[1],
|
|
dx[0],
|
|
dx[1]);
|
|
} else if (n == 3) {
|
|
const auto Ma = ins_dims[0][0];
|
|
const auto Ka = ins_dims[0][1];
|
|
const auto Nb = ins_dims[1][1];
|
|
const auto Nc = ins_dims[2][1];
|
|
const uint64_t cost1 = Ma * Nb * (Ka + Nc);
|
|
const uint64_t cost2 = Ka * Nc * (Nb + Ma);
|
|
auto mat_dim_a = funcs::CreateMatrixDescriptor(ins_dims[0], 0, false);
|
|
auto mat_dim_b = funcs::CreateMatrixDescriptor(ins_dims[1], 0, false);
|
|
auto mat_dim_c = funcs::CreateMatrixDescriptor(ins_dims[2], 0, false);
|
|
if (cost1 < cost2) {
|
|
DenseTensor tmp_out, tmp_dout;
|
|
tmp_out.Resize({Ma, Nb});
|
|
dev_ctx.template Alloc<T>(&tmp_out);
|
|
tmp_dout.Resize({mat_dim_dout.height_, Nb});
|
|
dev_ctx.template Alloc<T>(&tmp_dout);
|
|
blas.MatMul(
|
|
*ins[0], mat_dim_a, *ins[1], mat_dim_b, alpha, &tmp_out, T(0));
|
|
CalcGrad<Context, T>(dev_ctx,
|
|
dout,
|
|
tmp_out,
|
|
*ins[2],
|
|
dout_dim,
|
|
tmp_out.dims(),
|
|
ins_dims[2],
|
|
&tmp_dout,
|
|
dx[2]);
|
|
CalcGrad<Context, T>(dev_ctx,
|
|
tmp_dout,
|
|
*ins[0],
|
|
*ins[1],
|
|
tmp_dout.dims(),
|
|
ins_dims[0],
|
|
ins_dims[1],
|
|
dx[0],
|
|
dx[1]);
|
|
} else {
|
|
DenseTensor tmp_out, tmp_dout;
|
|
tmp_out.Resize({Ka, Nc});
|
|
dev_ctx.template Alloc<T>(&tmp_out);
|
|
tmp_dout.Resize({Ka, mat_dim_dout.width_});
|
|
dev_ctx.template Alloc<T>(&tmp_dout);
|
|
blas.MatMul(
|
|
*ins[1], mat_dim_b, *ins[2], mat_dim_c, alpha, &tmp_out, T(0));
|
|
CalcGrad<Context, T>(dev_ctx,
|
|
dout,
|
|
*ins[0],
|
|
tmp_out,
|
|
dout_dim,
|
|
ins_dims[0],
|
|
tmp_dout.dims(),
|
|
dx[0],
|
|
&tmp_dout);
|
|
CalcGrad<Context, T>(dev_ctx,
|
|
tmp_dout,
|
|
*ins[1],
|
|
*ins[2],
|
|
tmp_dout.dims(),
|
|
ins_dims[1],
|
|
ins_dims[2],
|
|
dx[1],
|
|
dx[2]);
|
|
}
|
|
} else {
|
|
MultiDotGradMatChainOrder<Context, T>(
|
|
dev_ctx, dout, ins, dout_dim, ins_dims, &dx);
|
|
// if x's shape is: [3] [3, 4] [4]
|
|
// dx's shape will be: [1, 3] [3, 4] [4, 1]
|
|
if (ins[n - 1]->dims().size() == 1) {
|
|
dx[n - 1]->Resize({dx[n - 1]->dims()[0]});
|
|
}
|
|
if (ins[0]->dims().size() == 1) {
|
|
dx[0]->Resize({dx[0]->dims()[1]});
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|