wehub-resource-sync
317 lines
11 KiB
C++
317 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.
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <complex>
|
|
|
|
#include "paddle/phi/backends/cpu/cpu_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/complex_functors.h"
|
|
#include "paddle/phi/kernels/funcs/lapack/lapack_function.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
#include "paddle/phi/kernels/impl/lstsq_kernel_impl.h"
|
|
#include "paddle/phi/kernels/lstsq_kernel.h"
|
|
#include "paddle/phi/kernels/transpose_kernel.h"
|
|
namespace phi {
|
|
|
|
enum class LapackDriverType : int { Gels, Gelsd, Gelsy, Gelss };
|
|
|
|
template <typename T, typename Context>
|
|
void LstsqKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& y,
|
|
const Scalar& rcond_scaler,
|
|
const std::string& driver_string,
|
|
DenseTensor* solution,
|
|
DenseTensor* residuals,
|
|
DenseTensor* rank,
|
|
DenseTensor* singular_values) {
|
|
using ValueType = dtype::Real<T>;
|
|
if (x.numel() == 0 || y.numel() == 0) {
|
|
if (solution) Full<T, Context>(dev_ctx, solution->dims(), 0, solution);
|
|
if (rank) Full<int64_t, Context>(dev_ctx, rank->dims(), 0, rank);
|
|
if (residuals)
|
|
GetResidualsTensor<Context, T>(
|
|
dev_ctx, x, y, driver_string, solution, residuals, rank);
|
|
if (singular_values)
|
|
Full<T, Context>(dev_ctx, singular_values->dims(), 0, singular_values);
|
|
return;
|
|
}
|
|
|
|
static auto driver_type = std::unordered_map<std::string, LapackDriverType>(
|
|
{{"gels", LapackDriverType::Gels},
|
|
{"gelsy", LapackDriverType::Gelsy},
|
|
{"gelsd", LapackDriverType::Gelsd},
|
|
{"gelss", LapackDriverType::Gelss}});
|
|
auto driver = driver_type[driver_string];
|
|
T rcond = rcond_scaler.to<T>();
|
|
|
|
auto x_dims = x.dims();
|
|
auto y_dims = y.dims();
|
|
int dim_size = x_dims.size();
|
|
int x_stride = GetMatrixStride(x_dims);
|
|
int y_stride = GetMatrixStride(y_dims);
|
|
int batch_count = GetBatchCount(x_dims);
|
|
auto solution_dim = solution->dims();
|
|
int ori_solu_stride = GetMatrixStride(solution_dim);
|
|
int max_solu_stride = std::max(y_stride, ori_solu_stride);
|
|
int min_solu_stride = std::min(y_stride, ori_solu_stride);
|
|
|
|
// lapack is a column-major storage, transpose make the input to
|
|
// have a continuous memory layout
|
|
int info = 0;
|
|
int m = static_cast<int>(x_dims[dim_size - 2]);
|
|
int n = static_cast<int>(x_dims[dim_size - 1]);
|
|
int nrhs = static_cast<int>(y_dims[dim_size - 1]);
|
|
int lda = std::max<int>(m, 1);
|
|
int ldb = std::max<int>(1, std::max(m, n));
|
|
|
|
DenseTensor new_x;
|
|
new_x.Resize({batch_count, m, n});
|
|
dev_ctx.template Alloc<T>(&new_x);
|
|
Copy<Context>(dev_ctx, x, dev_ctx.GetPlace(), true, &new_x);
|
|
|
|
solution->Resize({batch_count, std::max(m, n), nrhs});
|
|
dev_ctx.template Alloc<T>(solution);
|
|
|
|
if (m >= n) {
|
|
Copy<Context>(dev_ctx, y, dev_ctx.GetPlace(), true, solution);
|
|
} else {
|
|
auto* solu_data = solution->data<T>();
|
|
auto* y_data = y.data<T>();
|
|
for (auto i = 0; i < batch_count; i++) {
|
|
for (auto j = 0; j < min_solu_stride; j++) {
|
|
solu_data[i * max_solu_stride + j] = y_data[i * y_stride + j];
|
|
}
|
|
}
|
|
}
|
|
|
|
DenseTensor input_x_trans = TransposeLast2Dim<T>(dev_ctx, new_x);
|
|
DenseTensor input_y_trans = TransposeLast2Dim<T>(dev_ctx, *solution);
|
|
Copy<Context>(dev_ctx, input_x_trans, dev_ctx.GetPlace(), true, &new_x);
|
|
Copy<Context>(dev_ctx, input_y_trans, dev_ctx.GetPlace(), true, solution);
|
|
|
|
auto* x_vector = new_x.data<T>();
|
|
auto* y_vector = solution->data<T>();
|
|
|
|
// "gels" divers does not need to compute rank
|
|
int rank_32 = 0;
|
|
int* rank_data = nullptr;
|
|
int* rank_working_ptr = nullptr;
|
|
if (driver != LapackDriverType::Gels) {
|
|
rank_data = dev_ctx.template Alloc<int>(rank);
|
|
rank_working_ptr = rank_data;
|
|
}
|
|
|
|
// "gelsd" and "gelss" divers need to compute singular values
|
|
ValueType* s_data = nullptr;
|
|
ValueType* s_working_ptr = nullptr;
|
|
int s_stride = 0;
|
|
if (driver == LapackDriverType::Gelsd || driver == LapackDriverType::Gelss) {
|
|
s_data = dev_ctx.template Alloc<T>(singular_values);
|
|
s_working_ptr = s_data;
|
|
auto s_dims = singular_values->dims();
|
|
s_stride = static_cast<int>(s_dims[s_dims.size() - 1]);
|
|
}
|
|
|
|
// "jpvt" is only used for "gelsy" driver
|
|
DenseTensor jpvt;
|
|
int* jpvt_data = nullptr;
|
|
if (driver == LapackDriverType::Gelsy) {
|
|
jpvt.Resize({std::max<int>(1, n)});
|
|
jpvt_data = dev_ctx.template Alloc<int>(&jpvt);
|
|
}
|
|
|
|
// run once the driver, first to get the optimal workspace size
|
|
int lwork = -1;
|
|
T wkopt = 0.0;
|
|
ValueType rwkopt;
|
|
int iwkopt = 0;
|
|
|
|
if (driver == LapackDriverType::Gels) {
|
|
funcs::lapackGels(
|
|
'N', m, n, nrhs, x_vector, lda, y_vector, ldb, &wkopt, lwork, &info);
|
|
} else if (driver == LapackDriverType::Gelsd) {
|
|
funcs::lapackGelsd(m,
|
|
n,
|
|
nrhs,
|
|
x_vector,
|
|
lda,
|
|
y_vector,
|
|
ldb,
|
|
s_working_ptr,
|
|
static_cast<ValueType>(rcond),
|
|
&rank_32,
|
|
&wkopt,
|
|
lwork,
|
|
&rwkopt,
|
|
&iwkopt,
|
|
&info);
|
|
} else if (driver == LapackDriverType::Gelsy) {
|
|
funcs::lapackGelsy(m,
|
|
n,
|
|
nrhs,
|
|
x_vector,
|
|
lda,
|
|
y_vector,
|
|
ldb,
|
|
jpvt_data,
|
|
static_cast<ValueType>(rcond),
|
|
&rank_32,
|
|
&wkopt,
|
|
lwork,
|
|
&rwkopt,
|
|
&info);
|
|
} else if (driver == LapackDriverType::Gelss) {
|
|
funcs::lapackGelss(m,
|
|
n,
|
|
nrhs,
|
|
x_vector,
|
|
lda,
|
|
y_vector,
|
|
ldb,
|
|
s_working_ptr,
|
|
static_cast<ValueType>(rcond),
|
|
&rank_32,
|
|
&wkopt,
|
|
lwork,
|
|
&rwkopt,
|
|
&info);
|
|
}
|
|
|
|
lwork = std::max<int>(1, static_cast<int>(dtype::Real<T>(wkopt)));
|
|
DenseTensor work;
|
|
work.Resize({lwork});
|
|
T* work_data = dev_ctx.template Alloc<T>(&work);
|
|
|
|
// "rwork" only used for complex inputs and "gelsy/gelsd/gelss" drivers
|
|
DenseTensor rwork;
|
|
ValueType* rwork_data = nullptr;
|
|
if (IsComplexDtype(x.dtype()) && driver != LapackDriverType::Gels) {
|
|
int rwork_len = 0;
|
|
if (driver == LapackDriverType::Gelsy) {
|
|
rwork_len = std::max<int>(1, 2 * n);
|
|
} else if (driver == LapackDriverType::Gelss) {
|
|
rwork_len = std::max<int>(1, 5 * std::min(m, n));
|
|
} else if (driver == LapackDriverType::Gelsd) {
|
|
rwork_len = std::max<int>(1, rwkopt);
|
|
}
|
|
rwork.Resize({rwork_len});
|
|
rwork_data = dev_ctx.template Alloc<ValueType>(&rwork);
|
|
}
|
|
|
|
// "iwork" workspace array is relevant only for "gelsd" driver
|
|
DenseTensor iwork;
|
|
int* iwork_data = nullptr;
|
|
if (driver == LapackDriverType::Gelsd) {
|
|
iwork.Resize({std::max<int>(1, iwkopt)});
|
|
iwork_data = dev_ctx.template Alloc<int>(&iwork);
|
|
}
|
|
|
|
for (auto i = 0; i < batch_count; ++i) {
|
|
auto* x_input = &x_vector[i * x_stride];
|
|
auto* y_input = &y_vector[i * max_solu_stride];
|
|
rank_working_ptr = rank_working_ptr ? &rank_data[i] : nullptr;
|
|
s_working_ptr = s_working_ptr ? &s_data[i * s_stride] : nullptr;
|
|
|
|
if (driver == LapackDriverType::Gels) {
|
|
funcs::lapackGels(
|
|
'N', m, n, nrhs, x_input, lda, y_input, ldb, work_data, lwork, &info);
|
|
} else if (driver == LapackDriverType::Gelsd) {
|
|
funcs::lapackGelsd(m,
|
|
n,
|
|
nrhs,
|
|
x_input,
|
|
lda,
|
|
y_input,
|
|
ldb,
|
|
s_working_ptr,
|
|
static_cast<ValueType>(rcond),
|
|
&rank_32,
|
|
work_data,
|
|
lwork,
|
|
rwork_data,
|
|
iwork_data,
|
|
&info);
|
|
} else if (driver == LapackDriverType::Gelsy) {
|
|
funcs::lapackGelsy(m,
|
|
n,
|
|
nrhs,
|
|
x_input,
|
|
lda,
|
|
y_input,
|
|
ldb,
|
|
jpvt_data,
|
|
static_cast<ValueType>(rcond),
|
|
&rank_32,
|
|
work_data,
|
|
lwork,
|
|
rwork_data,
|
|
&info);
|
|
} else if (driver == LapackDriverType::Gelss) {
|
|
funcs::lapackGelss(m,
|
|
n,
|
|
nrhs,
|
|
x_input,
|
|
lda,
|
|
y_input,
|
|
ldb,
|
|
s_working_ptr,
|
|
static_cast<ValueType>(rcond),
|
|
&rank_32,
|
|
work_data,
|
|
lwork,
|
|
rwork_data,
|
|
&info);
|
|
}
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
info,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"For batch [%d]: Lapack info is not zero but [%d]", i, info));
|
|
|
|
if (rank_working_ptr) *rank_working_ptr = static_cast<int>(rank_32);
|
|
}
|
|
|
|
DenseTensor tmp_s = TransposeLast2Dim<T>(dev_ctx, *solution);
|
|
Copy<Context>(dev_ctx, tmp_s, dev_ctx.GetPlace(), true, solution);
|
|
|
|
if (m > n) {
|
|
auto* solu_data = solution->data<T>();
|
|
for (auto i = 1; i < batch_count; i++) {
|
|
for (auto j = 0; j < min_solu_stride; j++) {
|
|
solu_data[i * min_solu_stride + j] = solu_data[i * max_solu_stride + j];
|
|
}
|
|
}
|
|
}
|
|
|
|
if (batch_count > 1) {
|
|
solution->Resize(solution_dim);
|
|
} else {
|
|
solution->Resize({n, nrhs});
|
|
}
|
|
|
|
GetResidualsTensor<Context, T>(
|
|
dev_ctx, x, y, driver_string, solution, residuals, rank);
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(lstsq, CPU, ALL_LAYOUT, phi::LstsqKernel, float, double) {
|
|
kernel->OutputAt(2).SetDataType(phi::DataType::INT32);
|
|
}
|