161 lines
6.2 KiB
C++
161 lines
6.2 KiB
C++
// Copyright (c) 2025 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_MAGMA
|
|
#include "paddle/phi/kernels/funcs/magma/magma_function.h"
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/common/errors.h"
|
|
#include "paddle/phi/backends/dynload/magma.h"
|
|
#include "paddle/phi/common/data_type.h"
|
|
|
|
namespace phi::funcs {
|
|
|
|
void magmaEnsureInit() {
|
|
static std::once_flag magma_once_flag;
|
|
|
|
std::call_once(magma_once_flag, []() {
|
|
magma_int_t info = dynload::magma_init();
|
|
PADDLE_ENFORCE_EQ(
|
|
info,
|
|
0,
|
|
phi::errors::External("magma_init failed, info code = %d,"
|
|
"please checkout this code in: "
|
|
"https://github.com/icl-utk-edu/magma/blob/"
|
|
"master/include/magma_types.h#L542",
|
|
info));
|
|
|
|
std::atexit([]() {
|
|
magma_int_t info = dynload::magma_finalize();
|
|
PADDLE_ENFORCE_EQ(
|
|
info,
|
|
0,
|
|
phi::errors::External("magma_finalize failed, info code = %d,"
|
|
"please checkout this code in: "
|
|
"https://github.com/icl-utk-edu/magma/blob/"
|
|
"master/include/magma_types.h#L542",
|
|
info));
|
|
});
|
|
});
|
|
}
|
|
|
|
// Eig
|
|
template <>
|
|
void magmaEig<double>(magma_vec_t jobvl,
|
|
magma_vec_t jobvr,
|
|
magma_int_t n,
|
|
double *a,
|
|
magma_int_t lda,
|
|
double *w,
|
|
double *vl,
|
|
magma_int_t ldvl,
|
|
double *vr,
|
|
magma_int_t ldvr,
|
|
double *work,
|
|
magma_int_t lwork,
|
|
double *rwork,
|
|
magma_int_t *info) {
|
|
double *wr = w;
|
|
double *wi = w + n;
|
|
(void)rwork; // unused
|
|
*info = dynload::magma_dgeev(
|
|
jobvl, jobvr, n, a, lda, wr, wi, vl, ldvl, vr, ldvr, work, lwork, info);
|
|
}
|
|
|
|
template <>
|
|
void magmaEig<float>(magma_vec_t jobvl,
|
|
magma_vec_t jobvr,
|
|
magma_int_t n,
|
|
float *a,
|
|
magma_int_t lda,
|
|
float *w,
|
|
float *vl,
|
|
magma_int_t ldvl,
|
|
float *vr,
|
|
magma_int_t ldvr,
|
|
float *work,
|
|
magma_int_t lwork,
|
|
float *rwork,
|
|
magma_int_t *info) {
|
|
float *wr = w;
|
|
float *wi = w + n;
|
|
(void)rwork; // unused
|
|
*info = dynload::magma_sgeev(
|
|
jobvl, jobvr, n, a, lda, wr, wi, vl, ldvl, vr, ldvr, work, lwork, info);
|
|
}
|
|
|
|
template <>
|
|
void magmaEig<phi::complex128, double>(magma_vec_t jobvl,
|
|
magma_vec_t jobvr,
|
|
magma_int_t n,
|
|
phi::complex128 *a,
|
|
magma_int_t lda,
|
|
phi::complex128 *w,
|
|
phi::complex128 *vl,
|
|
magma_int_t ldvl,
|
|
phi::complex128 *vr,
|
|
magma_int_t ldvr,
|
|
phi::complex128 *work,
|
|
magma_int_t lwork,
|
|
double *rwork,
|
|
magma_int_t *info) {
|
|
*info = dynload::magma_zgeev(jobvl,
|
|
jobvr,
|
|
n,
|
|
reinterpret_cast<magmaDoubleComplex *>(a),
|
|
lda,
|
|
reinterpret_cast<magmaDoubleComplex *>(w),
|
|
reinterpret_cast<magmaDoubleComplex *>(vl),
|
|
ldvl,
|
|
reinterpret_cast<magmaDoubleComplex *>(vr),
|
|
ldvr,
|
|
reinterpret_cast<magmaDoubleComplex *>(work),
|
|
lwork,
|
|
rwork,
|
|
info);
|
|
}
|
|
|
|
template <>
|
|
void magmaEig<phi::complex64, float>(magma_vec_t jobvl,
|
|
magma_vec_t jobvr,
|
|
magma_int_t n,
|
|
phi::complex64 *a,
|
|
magma_int_t lda,
|
|
phi::complex64 *w,
|
|
phi::complex64 *vl,
|
|
magma_int_t ldvl,
|
|
phi::complex64 *vr,
|
|
magma_int_t ldvr,
|
|
phi::complex64 *work,
|
|
magma_int_t lwork,
|
|
float *rwork,
|
|
magma_int_t *info) {
|
|
*info = dynload::magma_cgeev(jobvl,
|
|
jobvr,
|
|
n,
|
|
reinterpret_cast<magmaFloatComplex *>(a),
|
|
lda,
|
|
reinterpret_cast<magmaFloatComplex *>(w),
|
|
reinterpret_cast<magmaFloatComplex *>(vl),
|
|
ldvl,
|
|
reinterpret_cast<magmaFloatComplex *>(vr),
|
|
ldvr,
|
|
reinterpret_cast<magmaFloatComplex *>(work),
|
|
lwork,
|
|
rwork,
|
|
info);
|
|
}
|
|
|
|
} // namespace phi::funcs
|
|
#endif
|