wehub-resource-sync
182 lines
6.7 KiB
Plaintext
182 lines
6.7 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/funcs/lapack/lapack_function.h"
|
|
#include "paddle/phi/kernels/impl/cholesky_solve_kernel_impl.h"
|
|
|
|
namespace phi {
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
template <typename T>
|
|
void rocsolver_potrs(const solverHandle_t &handle,
|
|
rocblas_fill uplo,
|
|
int M,
|
|
int N,
|
|
T *Adata,
|
|
int lda,
|
|
T *Bdata,
|
|
int ldb);
|
|
|
|
using dtype::complex;
|
|
#define FUNC_WITH_TYPES(m) \
|
|
m(float, s, float) m(double, d, double) \
|
|
m(complex<float>, c, rocblas_float_complex) \
|
|
m(complex<double>, z, rocblas_double_complex)
|
|
|
|
#define POTRS_INSTANCE(T, C, CastType) \
|
|
template <> \
|
|
void rocsolver_potrs<T>(const solverHandle_t &handle, \
|
|
rocblas_fill uplo, \
|
|
int M, \
|
|
int N, \
|
|
T *Adata, \
|
|
int lda, \
|
|
T *Bdata, \
|
|
int ldb) { \
|
|
PADDLE_ENFORCE_GPU_SUCCESS( \
|
|
dynload::rocsolver_##C##potrs(handle, \
|
|
uplo, \
|
|
M, \
|
|
N, \
|
|
reinterpret_cast<CastType *>(Adata), \
|
|
lda, \
|
|
reinterpret_cast<CastType *>(Bdata), \
|
|
ldb)); \
|
|
}
|
|
|
|
FUNC_WITH_TYPES(POTRS_INSTANCE);
|
|
|
|
#else
|
|
template <typename T>
|
|
void cusolver_potrs(const solverHandle_t &handle,
|
|
cublasFillMode_t uplo,
|
|
int M,
|
|
int N,
|
|
T *Adata,
|
|
int lda,
|
|
T *Bdata,
|
|
int ldb,
|
|
int *devInfo);
|
|
|
|
template <>
|
|
void cusolver_potrs<float>(const solverHandle_t &handle,
|
|
cublasFillMode_t uplo,
|
|
int M,
|
|
int N,
|
|
float *Adata,
|
|
int lda,
|
|
float *Bdata,
|
|
int ldb,
|
|
int *devInfo) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnSpotrs(
|
|
handle, uplo, M, N, Adata, lda, Bdata, ldb, devInfo));
|
|
}
|
|
|
|
template <>
|
|
void cusolver_potrs<double>(const solverHandle_t &handle,
|
|
cublasFillMode_t uplo,
|
|
int M,
|
|
int N,
|
|
double *Adata,
|
|
int lda,
|
|
double *Bdata,
|
|
int ldb,
|
|
int *devInfo) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnDpotrs(
|
|
handle, uplo, M, N, Adata, lda, Bdata, ldb, devInfo));
|
|
}
|
|
|
|
template <>
|
|
void cusolver_potrs<complex64>(const solverHandle_t &handle,
|
|
cublasFillMode_t uplo,
|
|
int M,
|
|
int N,
|
|
complex64 *Adata,
|
|
int lda,
|
|
complex64 *Bdata,
|
|
int ldb,
|
|
int *devInfo) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnCpotrs(handle,
|
|
uplo,
|
|
M,
|
|
N,
|
|
reinterpret_cast<const cuComplex *>(Adata),
|
|
lda,
|
|
reinterpret_cast<cuComplex *>(Bdata),
|
|
ldb,
|
|
devInfo));
|
|
}
|
|
|
|
template <>
|
|
void cusolver_potrs<complex128>(const cusolverDnHandle_t &handle,
|
|
cublasFillMode_t uplo,
|
|
int M,
|
|
int N,
|
|
complex128 *Adata,
|
|
int lda,
|
|
complex128 *Bdata,
|
|
int ldb,
|
|
int *devInfo) {
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnZpotrs(
|
|
handle,
|
|
uplo,
|
|
M,
|
|
N,
|
|
reinterpret_cast<const cuDoubleComplex *>(Adata),
|
|
lda,
|
|
reinterpret_cast<cuDoubleComplex *>(Bdata),
|
|
ldb,
|
|
devInfo));
|
|
}
|
|
|
|
#endif // PADDLE_WITH_HIP
|
|
|
|
template <typename T>
|
|
class CholeskySolveFunctor<T, GPUContext> {
|
|
public:
|
|
void operator()(const GPUContext &dev_ctx,
|
|
bool upper,
|
|
int M,
|
|
int N,
|
|
T *Adata,
|
|
int lda,
|
|
T *Bdata,
|
|
int *devInfo) {
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
#ifdef PADDLE_WITH_HIP
|
|
rocblas_fill uplo = upper ? rocblas_fill_upper : rocblas_fill_lower;
|
|
rocsolver_potrs<T>(handle, uplo, M, N, Adata, lda, Bdata, lda);
|
|
#else
|
|
cublasFillMode_t uplo =
|
|
upper ? CUBLAS_FILL_MODE_UPPER : CUBLAS_FILL_MODE_LOWER;
|
|
cusolver_potrs<T>(handle, uplo, M, N, Adata, lda, Bdata, lda, devInfo);
|
|
#endif
|
|
}
|
|
};
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(
|
|
cholesky_solve, GPU, ALL_LAYOUT, phi::CholeskySolveKernel, float, double) {}
|