wehub-resource-sync
362 lines
13 KiB
Plaintext
362 lines
13 KiB
Plaintext
// 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.
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
#include "paddle/phi/backends/dynload/rocsolver.h"
|
|
#else
|
|
#include "paddle/phi/backends/dynload/cusolver.h"
|
|
#endif
|
|
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/kernels/impl/lu_kernel_impl.h"
|
|
#include "paddle/phi/kernels/lu_kernel.h"
|
|
|
|
namespace phi {
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
template <typename T>
|
|
void rocsolver_getrf(const rocblas_handle& handle,
|
|
int m,
|
|
int n,
|
|
T* a,
|
|
int lda,
|
|
int* ipiv,
|
|
int* info);
|
|
|
|
template <>
|
|
void rocsolver_getrf<float>(const rocblas_handle& handle,
|
|
int m,
|
|
int n,
|
|
float* a,
|
|
int lda,
|
|
int* ipiv,
|
|
int* info) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::rocsolver_sgetrf(handle, m, n, a, lda, ipiv, info));
|
|
}
|
|
|
|
template <>
|
|
void rocsolver_getrf<double>(const rocblas_handle& handle,
|
|
int m,
|
|
int n,
|
|
double* a,
|
|
int lda,
|
|
int* ipiv,
|
|
int* info) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::rocsolver_dgetrf(handle, m, n, a, lda, ipiv, info));
|
|
}
|
|
|
|
template <>
|
|
void rocsolver_getrf<dtype::complex<float>>(const rocblas_handle& handle,
|
|
int m,
|
|
int n,
|
|
dtype::complex<float>* a,
|
|
int lda,
|
|
int* ipiv,
|
|
int* info) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::rocsolver_cgetrf(handle,
|
|
m,
|
|
n,
|
|
reinterpret_cast<rocblas_float_complex*>(a),
|
|
lda,
|
|
ipiv,
|
|
info));
|
|
}
|
|
|
|
template <>
|
|
void rocsolver_getrf<dtype::complex<double>>(const rocblas_handle& handle,
|
|
int m,
|
|
int n,
|
|
dtype::complex<double>* a,
|
|
int lda,
|
|
int* ipiv,
|
|
int* info) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::rocsolver_zgetrf(handle,
|
|
m,
|
|
n,
|
|
reinterpret_cast<rocblas_double_complex*>(a),
|
|
lda,
|
|
ipiv,
|
|
info));
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void lu_decomposed_kernel(const Context& dev_ctx,
|
|
int m,
|
|
int n,
|
|
T* d_A,
|
|
int lda,
|
|
int* d_Ipiv,
|
|
int* d_info) {
|
|
// rocSOLVER's getrf does not require a workspace buffer
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
rocsolver_getrf<T>(handle, m, n, d_A, lda, d_Ipiv, d_info);
|
|
PADDLE_ENFORCE_GPU_SUCCESS(hipDeviceSynchronize());
|
|
}
|
|
|
|
#else // PADDLE_WITH_CUDA
|
|
template <typename T>
|
|
void cusolver_bufferSize(const cusolverDnHandle_t& cusolverH,
|
|
int m,
|
|
int n,
|
|
T* d_A,
|
|
int lda,
|
|
int* lwork);
|
|
template <typename T>
|
|
void cusolver_getrf(const cusolverDnHandle_t& cusolverH,
|
|
int m,
|
|
int n,
|
|
T* d_A,
|
|
int lda,
|
|
T* d_work,
|
|
int* d_Ipiv,
|
|
int* d_info);
|
|
|
|
template <>
|
|
void cusolver_bufferSize<float>(const cusolverDnHandle_t& cusolverH,
|
|
int m,
|
|
int n,
|
|
float* d_A,
|
|
int lda,
|
|
int* lwork) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnSgetrf_bufferSize(cusolverH, m, n, d_A, lda, lwork));
|
|
}
|
|
|
|
template <>
|
|
void cusolver_bufferSize<double>(const cusolverDnHandle_t& cusolverH,
|
|
int m,
|
|
int n,
|
|
double* d_A,
|
|
int lda,
|
|
int* lwork) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnDgetrf_bufferSize(cusolverH, m, n, d_A, lda, lwork));
|
|
}
|
|
|
|
template <>
|
|
void cusolver_bufferSize<dtype::complex<float>>(
|
|
const cusolverDnHandle_t& cusolverH,
|
|
int m,
|
|
int n,
|
|
dtype::complex<float>* d_A,
|
|
int lda,
|
|
int* lwork) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnCgetrf_bufferSize(
|
|
cusolverH, m, n, reinterpret_cast<cuComplex*>(d_A), lda, lwork));
|
|
}
|
|
|
|
template <>
|
|
void cusolver_bufferSize<dtype::complex<double>>(
|
|
const cusolverDnHandle_t& cusolverH,
|
|
int m,
|
|
int n,
|
|
dtype::complex<double>* d_A,
|
|
int lda,
|
|
int* lwork) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnZgetrf_bufferSize(
|
|
cusolverH, m, n, reinterpret_cast<cuDoubleComplex*>(d_A), lda, lwork));
|
|
}
|
|
|
|
template <>
|
|
void cusolver_getrf<float>(const cusolverDnHandle_t& cusolverH,
|
|
int m,
|
|
int n,
|
|
float* d_A,
|
|
int lda,
|
|
float* d_work,
|
|
int* d_Ipiv,
|
|
int* d_info) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnSgetrf(
|
|
cusolverH, m, n, d_A, lda, d_work, d_Ipiv, d_info));
|
|
}
|
|
|
|
template <>
|
|
void cusolver_getrf<double>(const cusolverDnHandle_t& cusolverH,
|
|
int m,
|
|
int n,
|
|
double* d_A,
|
|
int lda,
|
|
double* d_work,
|
|
int* d_Ipiv,
|
|
int* d_info) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnDgetrf(
|
|
cusolverH, m, n, d_A, lda, d_work, d_Ipiv, d_info));
|
|
}
|
|
|
|
template <>
|
|
void cusolver_getrf<dtype::complex<float>>(const cusolverDnHandle_t& cusolverH,
|
|
int m,
|
|
int n,
|
|
dtype::complex<float>* d_A,
|
|
int lda,
|
|
dtype::complex<float>* d_work,
|
|
int* d_Ipiv,
|
|
int* d_info) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnCgetrf(cusolverH,
|
|
m,
|
|
n,
|
|
reinterpret_cast<cuComplex*>(d_A),
|
|
lda,
|
|
reinterpret_cast<cuComplex*>(d_work),
|
|
d_Ipiv,
|
|
d_info));
|
|
}
|
|
|
|
template <>
|
|
void cusolver_getrf<dtype::complex<double>>(const cusolverDnHandle_t& cusolverH,
|
|
int m,
|
|
int n,
|
|
dtype::complex<double>* d_A,
|
|
int lda,
|
|
dtype::complex<double>* d_work,
|
|
int* d_Ipiv,
|
|
int* d_info) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnZgetrf(cusolverH,
|
|
m,
|
|
n,
|
|
reinterpret_cast<cuDoubleComplex*>(d_A),
|
|
lda,
|
|
reinterpret_cast<cuDoubleComplex*>(d_work),
|
|
d_Ipiv,
|
|
d_info));
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void lu_decomposed_kernel(const Context& dev_ctx,
|
|
int m,
|
|
int n,
|
|
T* d_A,
|
|
int lda,
|
|
int* d_Ipiv,
|
|
int* d_info) {
|
|
/* step 1: get cusolver handle*/
|
|
auto cusolverH = dev_ctx.cusolver_dn_handle();
|
|
|
|
/* step 2: query working space of getrf */
|
|
int lwork;
|
|
cusolver_bufferSize(cusolverH, m, n, d_A, lda, &lwork);
|
|
|
|
auto work_buff =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(),
|
|
lwork * sizeof(T),
|
|
Stream(reinterpret_cast<StreamId>(dev_ctx.stream())));
|
|
T* d_work = reinterpret_cast<T*>(work_buff->ptr());
|
|
|
|
/* step 3: LU factorization */
|
|
if (d_Ipiv) {
|
|
cusolver_getrf(cusolverH, m, n, d_A, lda, d_work, d_Ipiv, d_info);
|
|
} else {
|
|
cusolver_getrf(cusolverH, m, n, d_A, lda, d_work, NULL, d_info);
|
|
}
|
|
PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize());
|
|
}
|
|
#endif
|
|
|
|
template <typename T, typename Context>
|
|
void LUKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
bool pivot,
|
|
DenseTensor* out,
|
|
DenseTensor* pivots,
|
|
DenseTensor* infos) {
|
|
// big tensor currently not supported
|
|
PADDLE_ENFORCE_GE(
|
|
x.dims().size(),
|
|
2,
|
|
::common::errors::PreconditionNotMet(
|
|
"Invalid input x dimensionality: %d (expected ≥2)", x.dims().size()));
|
|
if (x.numel() == 0) {
|
|
Full<int, Context>(dev_ctx, infos->dims(), static_cast<int>(0), infos);
|
|
Full<int, Context>(dev_ctx, pivots->dims(), static_cast<int>(0), pivots);
|
|
Full<T, Context>(dev_ctx, out->dims(), static_cast<T>(0), out);
|
|
return;
|
|
}
|
|
int64_t largest_matrix = (1LL << 31) - 1;
|
|
int64_t last = x.dims()[x.dims().size() - 1],
|
|
second_last = x.dims()[x.dims().size() - 2];
|
|
int64_t matrix_size = last * second_last;
|
|
PADDLE_ENFORCE_LE(matrix_size,
|
|
largest_matrix,
|
|
::common::errors::PreconditionNotMet(
|
|
"Matrix size too large for LU decomposition. Maximum "
|
|
"allowed size is 2 ^ 31 - 1 elements, but got %lld",
|
|
matrix_size));
|
|
|
|
const int64_t kMaxBlockDim = 512;
|
|
|
|
*out = Transpose2DTo6D<Context, T>(dev_ctx, x);
|
|
|
|
auto outdims = out->dims();
|
|
auto outrank = outdims.size();
|
|
|
|
int m = static_cast<int>(outdims[outrank - 1]);
|
|
int n = static_cast<int>(outdims[outrank - 2]);
|
|
int lda = std::max(1, m);
|
|
if (pivot) {
|
|
auto ipiv_dims = slice_ddim(outdims, 0, outrank - 1);
|
|
ipiv_dims[outrank - 2] = std::min(m, n);
|
|
pivots->Resize(ipiv_dims);
|
|
}
|
|
dev_ctx.template Alloc<int>(pivots);
|
|
auto ipiv_data = pivots->data<int>();
|
|
|
|
auto info_dims = slice_ddim(outdims, 0, outrank - 2);
|
|
infos->Resize(info_dims);
|
|
dev_ctx.template Alloc<int>(infos);
|
|
auto info_data = infos->data<int>();
|
|
|
|
auto batchsize = product(info_dims);
|
|
batchsize = std::max(static_cast<int>(batchsize), 1);
|
|
dev_ctx.template Alloc<T>(out);
|
|
auto out_data = out->data<T>();
|
|
for (int b = 0; b < batchsize; b++) {
|
|
auto out_data_item = &out_data[b * m * n];
|
|
int* info_data_item = &info_data[b];
|
|
if (pivot) {
|
|
auto ipiv_data_item = &ipiv_data[b * std::min(m, n)];
|
|
lu_decomposed_kernel(
|
|
dev_ctx, m, n, out_data_item, lda, ipiv_data_item, info_data_item);
|
|
} else {
|
|
lu_decomposed_kernel(
|
|
dev_ctx, m, n, out_data_item, lda, NULL, info_data_item);
|
|
}
|
|
}
|
|
*out = Transpose2DTo6D<Context, T>(dev_ctx, *out);
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(lu,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::LUKernel,
|
|
float,
|
|
double,
|
|
phi::complex64,
|
|
phi::complex128) {
|
|
kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
|
|
kernel->OutputAt(2).SetDataType(phi::DataType::INT32);
|
|
}
|