wehub-resource-sync
309 lines
11 KiB
C++
309 lines
11 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.
|
|
|
|
#include <fstream>
|
|
#include <iomanip>
|
|
#include <iostream>
|
|
#include <memory>
|
|
#include <vector>
|
|
#include "paddle/phi/backends/cpu/cpu_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
#include "paddle/phi/kernels/funcs/search_compute.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void CPUMatchMatrixTensorOPKernel(const Context& dev_ctx,
|
|
const DenseTensor& x_in,
|
|
const DenseTensor& y_in,
|
|
const DenseTensor& w_in,
|
|
int dim_t,
|
|
DenseTensor* out,
|
|
DenseTensor* tmp) {
|
|
auto* x = &x_in;
|
|
auto* y = &y_in;
|
|
auto* w = &w_in;
|
|
|
|
const auto& x_lod = x->lod();
|
|
PADDLE_ENFORCE_EQ(x_lod.empty(),
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"The Input(X) should hold LoD information, but "
|
|
"received Input(X).lod() is empty."));
|
|
const auto& x_lod_0 = x_lod[0];
|
|
PADDLE_ENFORCE_GE(x_lod_0.size(),
|
|
2,
|
|
common::errors::InvalidArgument(
|
|
"The dimensions of Input(X)'s LoD data should be "
|
|
"equal to 2, but received %d.",
|
|
x_lod_0.size()));
|
|
auto x_dims = x->dims();
|
|
PADDLE_ENFORCE_EQ(x_dims[0],
|
|
static_cast<int64_t>(x_lod_0.back()),
|
|
common::errors::InvalidArgument(
|
|
"The last element of Input(X)'s LoD data should be "
|
|
"equal to the first dimension of Input(X). "
|
|
"But received the last element of Input(X)'s LoD "
|
|
"data is %d, the first dimension of Input(X) is %d.",
|
|
x_lod_0.back(),
|
|
x_dims[0]));
|
|
const auto& y_lod = y->lod();
|
|
PADDLE_ENFORCE_EQ(y_lod.empty(),
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"The Input(Y) should hold LoD information, but "
|
|
"received Input(Y).lod() is empty."));
|
|
const auto& y_lod_0 = y_lod[0];
|
|
PADDLE_ENFORCE_GE(y_lod_0.size(),
|
|
2,
|
|
common::errors::InvalidArgument(
|
|
"The dimensions of Input(Y)'s LoD data should be "
|
|
"equal to 2, but received %d.",
|
|
y_lod_0.size()));
|
|
auto y_dims = y->dims();
|
|
PADDLE_ENFORCE_EQ(y_dims[0],
|
|
static_cast<int64_t>(y_lod_0.back()),
|
|
common::errors::InvalidArgument(
|
|
"The last element of Input(Y)'s LoD data should be "
|
|
"equal to the first dimension of Input(Y). "
|
|
"But received the last element of Input(Y)'s LoD "
|
|
"data is %d, the first dimension of Input(Y) is %d.",
|
|
y_lod_0.back(),
|
|
y_dims[0]));
|
|
|
|
PADDLE_ENFORCE_EQ(x_lod_0.size(),
|
|
y_lod_0.size(),
|
|
common::errors::InvalidArgument(
|
|
"The dimensions of Input(X)'s and Input(Y)'s LoD "
|
|
"data should be equal. "
|
|
"But received the dimensions of Input(X)'s LoD is "
|
|
"%d, the dimensions of Input(Y)'s LoD is %d.",
|
|
x_lod_0.size(),
|
|
y_lod_0.size()));
|
|
|
|
int64_t out_dim_0 = 0;
|
|
int64_t tmp_dim_0 = -1;
|
|
for (size_t i = 1; i < x_lod_0.size(); i++) {
|
|
int64_t x_len = x_lod_0[i] - x_lod_0[i - 1];
|
|
int64_t y_len = y_lod_0[i] - y_lod_0[i - 1];
|
|
out_dim_0 += (x_len * y_len);
|
|
}
|
|
out_dim_0 *= dim_t;
|
|
|
|
tmp_dim_0 = x_dims[0] * dim_t * x_dims[1];
|
|
std::vector<int64_t> out_dims_vec{out_dim_0};
|
|
out_dims_vec.push_back(1);
|
|
std::vector<int64_t> tmp_dims_vec{tmp_dim_0};
|
|
tmp_dims_vec.push_back(1);
|
|
|
|
auto& out_meta = out->meta();
|
|
DenseTensorMeta new_out_meta(out_meta.dtype,
|
|
make_ddim(out_dims_vec),
|
|
out_meta.layout,
|
|
out_meta.legacy_lod);
|
|
out->set_meta(new_out_meta);
|
|
|
|
auto& tmp_meta = tmp->meta();
|
|
DenseTensorMeta new_tmp_meta(tmp_meta.dtype,
|
|
make_ddim(tmp_dims_vec),
|
|
tmp_meta.layout,
|
|
tmp_meta.legacy_lod);
|
|
tmp->set_meta(new_tmp_meta);
|
|
|
|
int64_t dim_in = x->dims()[1];
|
|
|
|
const auto& offset_l = x->lod()[0];
|
|
const auto& offset_r = y->lod()[0];
|
|
|
|
std::vector<size_t> top_offset;
|
|
size_t top_size = 0;
|
|
top_offset.push_back(top_size);
|
|
for (size_t b = 0; b < x->lod()[0].size() - 1; b++) {
|
|
size_t len_l = offset_l[b + 1] - offset_l[b];
|
|
size_t len_r = offset_r[b + 1] - offset_r[b];
|
|
top_size += dim_t * len_l * len_r;
|
|
top_offset.push_back(top_size);
|
|
}
|
|
auto* out_data = dev_ctx.template Alloc<T>(out);
|
|
memset(out_data, 0.0, out->dims()[0] * out->dims()[1] * sizeof(T));
|
|
|
|
auto* bottom_l_data = x->data<T>();
|
|
auto* bottom_r_data = y->data<T>();
|
|
auto* t_data = w->data<T>();
|
|
auto* bottom_l_trans_data = dev_ctx.template Alloc<T>(tmp);
|
|
memset(bottom_l_trans_data, 0.0, tmp->dims()[0] * tmp->dims()[1] * sizeof(T));
|
|
|
|
auto blas = funcs::GetBlas<CPUContext, T>(dev_ctx);
|
|
|
|
funcs::call_gemm(blas,
|
|
CblasNoTrans,
|
|
CblasNoTrans,
|
|
x->dims()[0],
|
|
dim_t * dim_in,
|
|
dim_in,
|
|
1.0f,
|
|
bottom_l_data,
|
|
t_data,
|
|
0.0f,
|
|
bottom_l_trans_data);
|
|
|
|
for (size_t b = 0; b < x->lod()[0].size() - 1; b++) {
|
|
for (int t = 0; t < dim_t; t++) {
|
|
size_t len_l = offset_l[b + 1] - offset_l[b];
|
|
size_t len_r = offset_r[b + 1] - offset_r[b];
|
|
auto* top_data = out_data + top_offset[b] + t * len_l * len_r;
|
|
const auto* l_t_data =
|
|
bottom_l_trans_data + offset_l[b] * dim_t * dim_in + t * dim_in;
|
|
const auto* r_data = bottom_r_data + offset_r[b] * dim_in;
|
|
auto blas_2 = funcs::GetBlas<CPUContext, T>(dev_ctx);
|
|
funcs::call_gemm_with_lda(blas_2,
|
|
CblasNoTrans,
|
|
CblasTrans,
|
|
len_l,
|
|
len_r,
|
|
dim_in,
|
|
1.0f,
|
|
l_t_data,
|
|
r_data,
|
|
0.0f,
|
|
top_data,
|
|
dim_t * dim_in);
|
|
}
|
|
}
|
|
|
|
LegacyLoD out_lod;
|
|
out_lod.push_back(top_offset);
|
|
|
|
out->set_lod(out_lod);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void CPUMatchMatrixTensorOPGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x_in,
|
|
const DenseTensor& y_in,
|
|
const DenseTensor& w_in,
|
|
const DenseTensor& tmp_in,
|
|
const DenseTensor& out_grad,
|
|
int dim_t,
|
|
DenseTensor* x_grad,
|
|
DenseTensor* y_grad,
|
|
DenseTensor* w_grad) {
|
|
auto* x = &x_in;
|
|
auto* y = &y_in;
|
|
auto* w = &w_in;
|
|
auto* tmp = &tmp_in;
|
|
|
|
int64_t dim_in = x->dims()[1];
|
|
|
|
const auto& offset_l = x->lod()[0];
|
|
const auto& offset_r = y->lod()[0];
|
|
std::vector<size_t> top_offset;
|
|
size_t top_size = 0;
|
|
top_offset.push_back(top_size);
|
|
for (size_t b = 0; b < x->lod()[0].size() - 1; b++) {
|
|
size_t len_l = offset_l[b + 1] - offset_l[b];
|
|
size_t len_r = offset_r[b + 1] - offset_r[b];
|
|
top_size += dim_t * len_l * len_r;
|
|
top_offset.push_back(top_size);
|
|
}
|
|
|
|
auto* bottom_l_data = x->data<T>();
|
|
auto* bottom_r_data = y->data<T>();
|
|
auto* bottom_l_trans_data = tmp->data<T>();
|
|
|
|
auto* d_out = &out_grad;
|
|
auto* d_x = x_grad;
|
|
auto* d_y = y_grad;
|
|
|
|
DenseTensor tmp_grad;
|
|
tmp_grad.Resize(tmp->dims());
|
|
auto* d_tmp_data = dev_ctx.template Alloc<T>(&tmp_grad);
|
|
auto* top_diff = d_out->data<T>();
|
|
auto* bottom_l_diff = dev_ctx.template Alloc<T>(d_x);
|
|
auto* bottom_r_diff = dev_ctx.template Alloc<T>(d_y);
|
|
auto* bottom_l_trans_diff = const_cast<T*>(d_tmp_data);
|
|
memset(bottom_l_diff, 0.0, x->dims()[0] * x->dims()[1] * sizeof(T));
|
|
memset(bottom_r_diff, 0.0, y->dims()[0] * y->dims()[1] * sizeof(T));
|
|
memset(bottom_l_trans_diff, 0.0, tmp->dims()[0] * tmp->dims()[1] * sizeof(T));
|
|
|
|
for (size_t b = 0; b < x->lod()[0].size() - 1; b++) {
|
|
for (int t = 0; t < dim_t; t++) {
|
|
size_t len_l = offset_l[b + 1] - offset_l[b];
|
|
size_t len_r = offset_r[b + 1] - offset_r[b];
|
|
|
|
for (size_t i = 0; i < len_l; i++) {
|
|
for (size_t j = 0; j < len_r; j++) {
|
|
auto diff =
|
|
top_diff[top_offset[b] + t * len_l * len_r + i * len_r + j];
|
|
auto* l_trans_data = bottom_l_trans_data +
|
|
(offset_l[b] + i) * dim_in * dim_t + t * dim_in;
|
|
auto* l_trans_diff = bottom_l_trans_diff +
|
|
(offset_l[b] + i) * dim_in * dim_t + t * dim_in;
|
|
auto* r_data = bottom_r_data + (offset_r[b] + j) * dim_in;
|
|
auto* r_diff = bottom_r_diff + (offset_r[b] + j) * dim_in;
|
|
if (diff != 0.0) {
|
|
funcs::axpy(r_data, l_trans_diff, dim_in, diff);
|
|
funcs::axpy(l_trans_data, r_diff, dim_in, diff);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
auto blas = funcs::GetBlas<CPUContext, T>(dev_ctx);
|
|
|
|
auto* t_data = w->data<T>();
|
|
auto* d_w = w_grad;
|
|
auto* t_diff = dev_ctx.template Alloc<T>(d_w);
|
|
memset(t_diff, 0.0, w->dims()[0] * w->dims()[1] * w->dims()[2] * sizeof(T));
|
|
// bottom_diff
|
|
funcs::call_gemm(blas,
|
|
CblasNoTrans,
|
|
CblasTrans,
|
|
x->dims()[0],
|
|
dim_in,
|
|
dim_t * dim_in,
|
|
1.0f,
|
|
bottom_l_trans_diff,
|
|
t_data,
|
|
1.0f,
|
|
bottom_l_diff);
|
|
|
|
// t_diff
|
|
funcs::call_gemm(blas,
|
|
CblasTrans,
|
|
CblasNoTrans,
|
|
dim_in,
|
|
dim_t * dim_in,
|
|
x->dims()[0],
|
|
1.0f,
|
|
bottom_l_data,
|
|
bottom_l_trans_diff,
|
|
1.0f,
|
|
t_diff);
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(match_matrix_tensor,
|
|
CPU,
|
|
ALL_LAYOUT,
|
|
phi::CPUMatchMatrixTensorOPKernel,
|
|
float) {}
|
|
PD_REGISTER_KERNEL(match_matrix_tensor_grad,
|
|
CPU,
|
|
ALL_LAYOUT,
|
|
phi::CPUMatchMatrixTensorOPGradKernel,
|
|
float) {}
|