wehub-resource-sync
906 lines
32 KiB
Plaintext
906 lines
32 KiB
Plaintext
// 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/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
|
|
#ifdef PADDLE_WITH_CUDA
|
|
#include "paddle/phi/backends/dynload/cublas.h"
|
|
#include "paddle/phi/backends/dynload/cusolver.h"
|
|
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
|
|
#endif // PADDLE_WITH_CUDA
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
#include "hip/hip_runtime.h"
|
|
#include "paddle/phi/backends/dynload/rocblas.h"
|
|
#include "paddle/phi/backends/dynload/rocsolver.h"
|
|
#include "rocblas/rocblas.h"
|
|
#include "rocsolver/rocsolver.h"
|
|
#endif // PADDLE_WITH_HIP
|
|
|
|
#endif // PADDLE_WITH_MAGMA
|
|
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/complex_kernel.h"
|
|
#include "paddle/phi/kernels/cpu/eig.h"
|
|
#include "paddle/phi/kernels/funcs/complex_functors.h"
|
|
#include "paddle/phi/kernels/funcs/for_range.h"
|
|
#include "paddle/phi/kernels/funcs/unsqueeze.h"
|
|
#include "paddle/phi/kernels/transpose_kernel.h"
|
|
|
|
namespace phi {
|
|
|
|
#ifdef PADDLE_WITH_MAGMA
|
|
template <typename T>
|
|
void SolveLinearSystemGPU(const GPUContext& dev_ctx,
|
|
const T* matrix_data,
|
|
const T* rhs_data,
|
|
T* out_data,
|
|
int order,
|
|
int rhs_cols,
|
|
int batch_count);
|
|
|
|
#ifdef PADDLE_WITH_CUDA
|
|
template <>
|
|
void SolveLinearSystemGPU<dtype::complex<float>>(
|
|
const GPUContext& dev_ctx,
|
|
const dtype::complex<float>*
|
|
matrix_data, // device ptr, row-major, size batch*order*order
|
|
const dtype::complex<float>*
|
|
rhs_data, // device ptr, row-major, size batch*order*rhs_cols
|
|
dtype::complex<float>*
|
|
out_data, // device ptr, row-major, size batch*order*rhs_cols
|
|
int order,
|
|
int rhs_cols,
|
|
int batch_count) {
|
|
// handles
|
|
cublasHandle_t cublas_handle = dev_ctx.cublas_handle();
|
|
cusolverDnHandle_t cusolver_handle = dev_ctx.cusolver_dn_handle();
|
|
auto stream = Stream(reinterpret_cast<StreamId>(dev_ctx.stream()));
|
|
|
|
// cuComplex constants
|
|
const cuComplex kAlpha = make_cuFloatComplex(1.0f, 0.0f);
|
|
const cuComplex kZero = make_cuFloatComplex(0.0f, 0.0f);
|
|
|
|
// Sizes
|
|
const size_t A_one_bytes =
|
|
static_cast<size_t>(order) * order * sizeof(cuComplex);
|
|
const size_t B_one_bytes =
|
|
static_cast<size_t>(order) * rhs_cols * sizeof(cuComplex);
|
|
const size_t A_batch_bytes = A_one_bytes * batch_count;
|
|
const size_t B_batch_bytes = B_one_bytes * batch_count;
|
|
|
|
const cuComplex* A_row_all = reinterpret_cast<const cuComplex*>(matrix_data);
|
|
const cuComplex* B_row_all = reinterpret_cast<const cuComplex*>(rhs_data);
|
|
cuComplex* X_row_all = reinterpret_cast<cuComplex*>(out_data);
|
|
|
|
auto dA_col_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(), A_batch_bytes, stream);
|
|
auto dB_col_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(), B_batch_bytes, stream);
|
|
cuComplex* dA_col = reinterpret_cast<cuComplex*>(dA_col_alloc->ptr());
|
|
cuComplex* dB_col = reinterpret_cast<cuComplex*>(dB_col_alloc->ptr());
|
|
|
|
auto d_pivots_alloc = memory_utils::Alloc(
|
|
dev_ctx.GetPlace(),
|
|
static_cast<size_t>(batch_count) * order * sizeof(int),
|
|
stream);
|
|
int* d_pivots = reinterpret_cast<int*>(d_pivots_alloc->ptr());
|
|
|
|
auto d_info_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(),
|
|
static_cast<size_t>(batch_count) * sizeof(int),
|
|
stream);
|
|
int* d_info = reinterpret_cast<int*>(d_info_alloc->ptr());
|
|
|
|
// A_row layout: row-major (order x order), B_row layout: row-major (order
|
|
// x rhs_cols)
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
const cuComplex* A_row = A_row_all + static_cast<size_t>(i) * order * order;
|
|
cuComplex* A_col = dA_col + static_cast<size_t>(i) * order * order;
|
|
const cuComplex* B_row =
|
|
B_row_all + static_cast<size_t>(i) * order * rhs_cols;
|
|
cuComplex* B_col = dB_col + static_cast<size_t>(i) * order * rhs_cols;
|
|
|
|
// transpose A_row (row-major) -> A_col (column-major) via C = A^T
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cublasCgeam(cublas_handle,
|
|
CUBLAS_OP_T,
|
|
CUBLAS_OP_N,
|
|
order,
|
|
order,
|
|
&kAlpha,
|
|
A_row,
|
|
order, // lda: when interpreting A_row as (order x
|
|
// order) row-major, using order
|
|
&kZero,
|
|
nullptr,
|
|
order,
|
|
A_col,
|
|
order)); // ldc = order (column-major leading dim)
|
|
|
|
// transpose B_row (row-major order x rhs_cols) -> B_col (column-major order
|
|
// x rhs_cols)
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasCgeam(
|
|
cublas_handle,
|
|
CUBLAS_OP_T,
|
|
CUBLAS_OP_N,
|
|
order,
|
|
rhs_cols,
|
|
&kAlpha,
|
|
B_row,
|
|
rhs_cols, // lda when A_row is viewed row-major: leading = rhs_cols
|
|
&kZero,
|
|
nullptr,
|
|
rhs_cols,
|
|
B_col,
|
|
order)); // ldc = order
|
|
}
|
|
|
|
int lwork = 0;
|
|
cuComplex* dA_col0 = dA_col;
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnCgetrf_bufferSize(
|
|
cusolver_handle, order, order, dA_col0, order, &lwork));
|
|
|
|
size_t work_bytes = static_cast<size_t>(lwork) * sizeof(cuComplex);
|
|
auto d_work_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(), work_bytes, stream);
|
|
cuComplex* d_work = reinterpret_cast<cuComplex*>(d_work_alloc->ptr());
|
|
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
cuComplex* A_col = dA_col + static_cast<size_t>(i) * order * order;
|
|
cuComplex* B_col = dB_col + static_cast<size_t>(i) * order * rhs_cols;
|
|
int* pivots_i = d_pivots + static_cast<size_t>(i) * order;
|
|
int* info_i = d_info + i;
|
|
|
|
// getrf (LU factorization) on A_col (column-major)
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnCgetrf(
|
|
cusolver_handle, order, order, A_col, order, d_work, pivots_i, info_i));
|
|
|
|
// getrs: solve A_col * X_col = B_col
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnCgetrs(
|
|
cusolver_handle,
|
|
CUBLAS_OP_N, // no transpose on column-major matrix
|
|
order,
|
|
rhs_cols,
|
|
A_col,
|
|
order,
|
|
pivots_i,
|
|
B_col,
|
|
order,
|
|
info_i));
|
|
}
|
|
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
cuComplex* B_col = dB_col + static_cast<size_t>(i) * order *
|
|
rhs_cols; // X in column-major
|
|
cuComplex* X_row = X_row_all + static_cast<size_t>(i) * order *
|
|
rhs_cols; // target row-major
|
|
|
|
// transpose X_col -> X_row
|
|
// We use C = A^T : A has shape (order x rhs_cols) in column-major, so C
|
|
// will be (rhs_cols x order), but we want X_row with shape (order x
|
|
// rhs_cols) in row-major; calling cublasCgeam with op=T and adjusted dims
|
|
// works:
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasCgeam(
|
|
cublas_handle,
|
|
CUBLAS_OP_T,
|
|
CUBLAS_OP_N,
|
|
rhs_cols,
|
|
order, // rowsC = rhs_cols, colsC = order
|
|
&kAlpha,
|
|
B_col,
|
|
order, // B_col lda = order (col-major)
|
|
&kZero,
|
|
nullptr,
|
|
order,
|
|
X_row,
|
|
rhs_cols)); // X_row ldc = rhs_cols (row-major leading dimension)
|
|
}
|
|
|
|
std::vector<int> h_info(batch_count, 0);
|
|
PADDLE_ENFORCE_EQ(
|
|
backends::gpu::IsCUDAGraphCapturing(),
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"EigGradKernel does not support CUDA Graph capture: async D2H copy "
|
|
"to local vector 'h_info' will bake the destination address into the "
|
|
"graph; on replay the vector is re-created at a different address, "
|
|
"causing a dangling-pointer write."));
|
|
phi::memory_utils::Copy(CPUPlace(),
|
|
h_info.data(),
|
|
dev_ctx.GetPlace(),
|
|
d_info,
|
|
static_cast<size_t>(batch_count) * sizeof(int),
|
|
reinterpret_cast<void*>(dev_ctx.stream()));
|
|
dev_ctx.Wait();
|
|
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
PADDLE_ENFORCE_EQ(
|
|
h_info[i],
|
|
0,
|
|
errors::External(
|
|
"cuSOLVER getrf/getrs failed at batch %d, info: %d", i, h_info[i]));
|
|
}
|
|
}
|
|
|
|
template <>
|
|
void SolveLinearSystemGPU<dtype::complex<double>>(
|
|
const GPUContext& dev_ctx,
|
|
const dtype::complex<double>*
|
|
matrix_data, // device ptr, row-major, size batch*order*order
|
|
const dtype::complex<double>*
|
|
rhs_data, // device ptr, row-major, size batch*order*rhs_cols
|
|
dtype::complex<double>*
|
|
out_data, // device ptr, row-major, size batch*order*rhs_cols
|
|
int order,
|
|
int rhs_cols,
|
|
int batch_count) {
|
|
// handles
|
|
cublasHandle_t cublas_handle = dev_ctx.cublas_handle();
|
|
cusolverDnHandle_t cusolver_handle = dev_ctx.cusolver_dn_handle();
|
|
auto stream = Stream(reinterpret_cast<StreamId>(dev_ctx.stream()));
|
|
|
|
// cuDoubleComplex constants
|
|
const cuDoubleComplex kAlpha = make_cuDoubleComplex(1.0f, 0.0f);
|
|
const cuDoubleComplex kZero = make_cuDoubleComplex(0.0f, 0.0f);
|
|
|
|
// Sizes
|
|
const size_t A_one_bytes =
|
|
static_cast<size_t>(order) * order * sizeof(cuDoubleComplex);
|
|
const size_t B_one_bytes =
|
|
static_cast<size_t>(order) * rhs_cols * sizeof(cuDoubleComplex);
|
|
const size_t A_batch_bytes = A_one_bytes * batch_count;
|
|
const size_t B_batch_bytes = B_one_bytes * batch_count;
|
|
|
|
const cuDoubleComplex* A_row_all =
|
|
reinterpret_cast<const cuDoubleComplex*>(matrix_data);
|
|
const cuDoubleComplex* B_row_all =
|
|
reinterpret_cast<const cuDoubleComplex*>(rhs_data);
|
|
cuDoubleComplex* X_row_all = reinterpret_cast<cuDoubleComplex*>(out_data);
|
|
|
|
auto dA_col_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(), A_batch_bytes, stream);
|
|
auto dB_col_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(), B_batch_bytes, stream);
|
|
cuDoubleComplex* dA_col =
|
|
reinterpret_cast<cuDoubleComplex*>(dA_col_alloc->ptr());
|
|
cuDoubleComplex* dB_col =
|
|
reinterpret_cast<cuDoubleComplex*>(dB_col_alloc->ptr());
|
|
|
|
auto d_pivots_alloc = memory_utils::Alloc(
|
|
dev_ctx.GetPlace(),
|
|
static_cast<size_t>(batch_count) * order * sizeof(int),
|
|
stream);
|
|
int* d_pivots = reinterpret_cast<int*>(d_pivots_alloc->ptr());
|
|
|
|
auto d_info_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(),
|
|
static_cast<size_t>(batch_count) * sizeof(int),
|
|
stream);
|
|
int* d_info = reinterpret_cast<int*>(d_info_alloc->ptr());
|
|
|
|
// A_row layout: row-major (order x order), B_row layout: row-major (order
|
|
// x rhs_cols)
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
const cuDoubleComplex* A_row =
|
|
A_row_all + static_cast<size_t>(i) * order * order;
|
|
cuDoubleComplex* A_col = dA_col + static_cast<size_t>(i) * order * order;
|
|
const cuDoubleComplex* B_row =
|
|
B_row_all + static_cast<size_t>(i) * order * rhs_cols;
|
|
cuDoubleComplex* B_col = dB_col + static_cast<size_t>(i) * order * rhs_cols;
|
|
|
|
// transpose A_row (row-major) -> A_col (column-major) via C = A^T
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cublasZgeam(cublas_handle,
|
|
CUBLAS_OP_T,
|
|
CUBLAS_OP_N,
|
|
order,
|
|
order,
|
|
&kAlpha,
|
|
A_row,
|
|
order, // lda: when interpreting A_row as (order x
|
|
// order) row-major, using order
|
|
&kZero,
|
|
nullptr,
|
|
order,
|
|
A_col,
|
|
order)); // ldc = order (column-major leading dim)
|
|
|
|
// transpose B_row (row-major order x rhs_cols) -> B_col (column-major order
|
|
// x rhs_cols)
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasZgeam(
|
|
cublas_handle,
|
|
CUBLAS_OP_T,
|
|
CUBLAS_OP_N,
|
|
order,
|
|
rhs_cols,
|
|
&kAlpha,
|
|
B_row,
|
|
rhs_cols, // lda when A_row is viewed row-major: leading = rhs_cols
|
|
&kZero,
|
|
nullptr,
|
|
rhs_cols,
|
|
B_col,
|
|
order)); // ldc = order
|
|
}
|
|
|
|
int lwork = 0;
|
|
cuDoubleComplex* dA_col0 = dA_col;
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnZgetrf_bufferSize(
|
|
cusolver_handle, order, order, dA_col0, order, &lwork));
|
|
|
|
size_t work_bytes = static_cast<size_t>(lwork) * sizeof(cuDoubleComplex);
|
|
auto d_work_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(), work_bytes, stream);
|
|
cuDoubleComplex* d_work =
|
|
reinterpret_cast<cuDoubleComplex*>(d_work_alloc->ptr());
|
|
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
cuDoubleComplex* A_col = dA_col + static_cast<size_t>(i) * order * order;
|
|
cuDoubleComplex* B_col = dB_col + static_cast<size_t>(i) * order * rhs_cols;
|
|
int* pivots_i = d_pivots + static_cast<size_t>(i) * order;
|
|
int* info_i = d_info + i;
|
|
|
|
// getrf (LU factorization) on A_col (column-major)
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnZgetrf(
|
|
cusolver_handle, order, order, A_col, order, d_work, pivots_i, info_i));
|
|
|
|
// getrs: solve A_col * X_col = B_col
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnZgetrs(
|
|
cusolver_handle,
|
|
CUBLAS_OP_N, // no transpose on column-major matrix
|
|
order,
|
|
rhs_cols,
|
|
A_col,
|
|
order,
|
|
pivots_i,
|
|
B_col,
|
|
order,
|
|
info_i));
|
|
}
|
|
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
cuDoubleComplex* B_col = dB_col + static_cast<size_t>(i) * order *
|
|
rhs_cols; // X in column-major
|
|
cuDoubleComplex* X_row = X_row_all + static_cast<size_t>(i) * order *
|
|
rhs_cols; // target row-major
|
|
|
|
// transpose X_col -> X_row
|
|
// We use C = A^T : A has shape (order x rhs_cols) in column-major, so C
|
|
// will be (rhs_cols x order), but we want X_row with shape (order x
|
|
// rhs_cols) in row-major; calling cublasZgeam with op=T and adjusted dims
|
|
// works:
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasZgeam(
|
|
cublas_handle,
|
|
CUBLAS_OP_T,
|
|
CUBLAS_OP_N,
|
|
rhs_cols,
|
|
order, // rowsC = rhs_cols, colsC = order
|
|
&kAlpha,
|
|
B_col,
|
|
order, // B_col lda = order (col-major)
|
|
&kZero,
|
|
nullptr,
|
|
order,
|
|
X_row,
|
|
rhs_cols)); // X_row ldc = rhs_cols (row-major leading dimension)
|
|
}
|
|
|
|
std::vector<int> h_info(batch_count, 0);
|
|
PADDLE_ENFORCE_EQ(
|
|
backends::gpu::IsCUDAGraphCapturing(),
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"EigGradKernel does not support CUDA Graph capture: async D2H copy "
|
|
"to local vector 'h_info' will bake the destination address into the "
|
|
"graph; on replay the vector is re-created at a different address, "
|
|
"causing a dangling-pointer write."));
|
|
phi::memory_utils::Copy(CPUPlace(),
|
|
h_info.data(),
|
|
dev_ctx.GetPlace(),
|
|
d_info,
|
|
static_cast<size_t>(batch_count) * sizeof(int),
|
|
reinterpret_cast<void*>(dev_ctx.stream()));
|
|
dev_ctx.Wait();
|
|
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
PADDLE_ENFORCE_EQ(
|
|
h_info[i],
|
|
0,
|
|
errors::External(
|
|
"cuSOLVER getrf/getrs failed at batch %d, info: %d", i, h_info[i]));
|
|
}
|
|
}
|
|
#endif // PADDLE_WITH_CUDA
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
template <>
|
|
void SolveLinearSystemGPU<dtype::complex<float>>(
|
|
const GPUContext& dev_ctx,
|
|
const dtype::complex<float>*
|
|
matrix_data, // device ptr, row-major, size batch*order*order
|
|
const dtype::complex<float>*
|
|
rhs_data, // device ptr, row-major, size batch*order*rhs_cols
|
|
dtype::complex<float>*
|
|
out_data, // device ptr, row-major, size batch*order*rhs_cols
|
|
int order,
|
|
int rhs_cols,
|
|
int batch_count) {
|
|
// handles
|
|
rocblas_handle rocblas_handle = dev_ctx.cusolver_dn_handle();
|
|
auto stream = Stream(reinterpret_cast<StreamId>(dev_ctx.stream()));
|
|
|
|
// rocblas_float_complex constants
|
|
const rocblas_float_complex kAlpha = rocblas_float_complex{1.0f, 0.0f};
|
|
const rocblas_float_complex kZero = rocblas_float_complex{0.0f, 0.0f};
|
|
|
|
// Sizes
|
|
const size_t A_one_bytes =
|
|
static_cast<size_t>(order) * order * sizeof(rocblas_float_complex);
|
|
const size_t B_one_bytes =
|
|
static_cast<size_t>(order) * rhs_cols * sizeof(rocblas_float_complex);
|
|
const size_t A_batch_bytes = A_one_bytes * batch_count;
|
|
const size_t B_batch_bytes = B_one_bytes * batch_count;
|
|
|
|
const rocblas_float_complex* A_row_all =
|
|
reinterpret_cast<const rocblas_float_complex*>(matrix_data);
|
|
const rocblas_float_complex* B_row_all =
|
|
reinterpret_cast<const rocblas_float_complex*>(rhs_data);
|
|
rocblas_float_complex* X_row_all =
|
|
reinterpret_cast<rocblas_float_complex*>(out_data);
|
|
|
|
auto dA_col_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(), A_batch_bytes, stream);
|
|
auto dB_col_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(), B_batch_bytes, stream);
|
|
rocblas_float_complex* dA_col =
|
|
reinterpret_cast<rocblas_float_complex*>(dA_col_alloc->ptr());
|
|
rocblas_float_complex* dB_col =
|
|
reinterpret_cast<rocblas_float_complex*>(dB_col_alloc->ptr());
|
|
|
|
auto d_pivots_alloc = memory_utils::Alloc(
|
|
dev_ctx.GetPlace(),
|
|
static_cast<size_t>(batch_count) * order * sizeof(rocblas_int),
|
|
stream);
|
|
rocblas_int* d_pivots = reinterpret_cast<rocblas_int*>(d_pivots_alloc->ptr());
|
|
|
|
auto d_info_alloc = memory_utils::Alloc(
|
|
dev_ctx.GetPlace(),
|
|
static_cast<size_t>(batch_count) * sizeof(rocblas_int),
|
|
stream);
|
|
rocblas_int* d_info = reinterpret_cast<rocblas_int*>(d_info_alloc->ptr());
|
|
|
|
// A_row layout: row-major (order x order), B_row layout: row-major (order x
|
|
// rhs_cols)
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
const rocblas_float_complex* A_row =
|
|
A_row_all + static_cast<size_t>(i) * order * order;
|
|
rocblas_float_complex* A_col =
|
|
dA_col + static_cast<size_t>(i) * order * order;
|
|
const rocblas_float_complex* B_row =
|
|
B_row_all + static_cast<size_t>(i) * order * rhs_cols;
|
|
rocblas_float_complex* B_col =
|
|
dB_col + static_cast<size_t>(i) * order * rhs_cols;
|
|
|
|
// transpose A_row (row-major) -> A_col (column-major) via C = A^T
|
|
PADDLE_ENFORCE_GPU_SUCCESS(rocblas_cgeam(
|
|
rocblas_handle,
|
|
rocblas_operation_transpose,
|
|
rocblas_operation_none,
|
|
order,
|
|
order,
|
|
&kAlpha,
|
|
A_row,
|
|
order, // lda: when interpreting A_row as (order x order) row-major
|
|
&kZero,
|
|
nullptr,
|
|
order,
|
|
A_col,
|
|
order)); // ldc = order (column-major leading dim)
|
|
|
|
// transpose B_row (row-major order x rhs_cols) -> B_col (column-major order
|
|
// x rhs_cols)
|
|
PADDLE_ENFORCE_GPU_SUCCESS(rocblas_cgeam(
|
|
rocblas_handle,
|
|
rocblas_operation_transpose,
|
|
rocblas_operation_none,
|
|
order,
|
|
rhs_cols,
|
|
&kAlpha,
|
|
B_row,
|
|
rhs_cols, // lda when A_row is viewed row-major: leading = rhs_cols
|
|
&kZero,
|
|
nullptr,
|
|
rhs_cols,
|
|
B_col,
|
|
order)); // ldc = order
|
|
}
|
|
|
|
// LU factorization and solve for each batch
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
rocblas_float_complex* A_col =
|
|
dA_col + static_cast<size_t>(i) * order * order;
|
|
rocblas_float_complex* B_col =
|
|
dB_col + static_cast<size_t>(i) * order * rhs_cols;
|
|
rocblas_int* pivots_i = d_pivots + static_cast<size_t>(i) * order;
|
|
rocblas_int* info_i = d_info + i;
|
|
|
|
// getrf (LU factorization) on A_col (column-major)
|
|
PADDLE_ENFORCE_GPU_SUCCESS(rocsolver_cgetrf(
|
|
rocblas_handle, order, order, A_col, order, pivots_i, info_i));
|
|
|
|
// getrs: solve A_col * X_col = B_col
|
|
PADDLE_ENFORCE_GPU_SUCCESS(rocsolver_cgetrs(
|
|
rocblas_handle,
|
|
rocblas_operation_none, // no transpose on column-major matrix
|
|
order,
|
|
rhs_cols,
|
|
A_col,
|
|
order,
|
|
pivots_i,
|
|
B_col,
|
|
order));
|
|
}
|
|
|
|
// Transpose results back to row-major
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
rocblas_float_complex* B_col = dB_col + static_cast<size_t>(i) * order *
|
|
rhs_cols; // X in column-major
|
|
rocblas_float_complex* X_row =
|
|
X_row_all +
|
|
static_cast<size_t>(i) * order * rhs_cols; // target row-major
|
|
|
|
// transpose X_col -> X_row
|
|
PADDLE_ENFORCE_GPU_SUCCESS(rocblas_cgeam(
|
|
rocblas_handle,
|
|
rocblas_operation_transpose,
|
|
rocblas_operation_none,
|
|
rhs_cols,
|
|
order, // rowsC = rhs_cols, colsC = order
|
|
&kAlpha,
|
|
B_col,
|
|
order, // B_col lda = order (col-major)
|
|
&kZero,
|
|
nullptr,
|
|
order,
|
|
X_row,
|
|
rhs_cols)); // X_row ldc = rhs_cols (row-major leading dimension)
|
|
}
|
|
|
|
// Check error info
|
|
CPUPlace cpu_place;
|
|
DeviceContextPool& pool = DeviceContextPool::Instance();
|
|
auto* cpu_ctx = static_cast<CPUContext*>(pool.Get(cpu_place));
|
|
|
|
std::vector<rocblas_int> h_info(batch_count, 0);
|
|
memory_utils::Copy(CPUPlace(),
|
|
h_info.data(),
|
|
dev_ctx.GetPlace(),
|
|
d_info,
|
|
static_cast<size_t>(batch_count) * sizeof(rocblas_int),
|
|
reinterpret_cast<void*>(dev_ctx.stream()));
|
|
dev_ctx.Wait();
|
|
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
PADDLE_ENFORCE_EQ(
|
|
h_info[i],
|
|
0,
|
|
errors::External("rocSOLVER getrf/getrs failed at batch %d, info: %d",
|
|
i,
|
|
h_info[i]));
|
|
}
|
|
}
|
|
|
|
template <>
|
|
void SolveLinearSystemGPU<dtype::complex<double>>(
|
|
const GPUContext& dev_ctx,
|
|
const dtype::complex<double>*
|
|
matrix_data, // device ptr, row-major, size batch*order*order
|
|
const dtype::complex<double>*
|
|
rhs_data, // device ptr, row-major, size batch*order*rhs_cols
|
|
dtype::complex<double>*
|
|
out_data, // device ptr, row-major, size batch*order*rhs_cols
|
|
int order,
|
|
int rhs_cols,
|
|
int batch_count) {
|
|
// handles
|
|
rocblas_handle rocblas_handle = dev_ctx.cusolver_dn_handle();
|
|
auto stream = Stream(reinterpret_cast<StreamId>(dev_ctx.stream()));
|
|
|
|
// rocblas_double_complex constants
|
|
const rocblas_double_complex kAlpha = rocblas_double_complex{1.0, 0.0};
|
|
const rocblas_double_complex kZero = rocblas_double_complex{0.0, 0.0};
|
|
|
|
// Sizes
|
|
const size_t A_one_bytes =
|
|
static_cast<size_t>(order) * order * sizeof(rocblas_double_complex);
|
|
const size_t B_one_bytes =
|
|
static_cast<size_t>(order) * rhs_cols * sizeof(rocblas_double_complex);
|
|
const size_t A_batch_bytes = A_one_bytes * batch_count;
|
|
const size_t B_batch_bytes = B_one_bytes * batch_count;
|
|
|
|
const rocblas_double_complex* A_row_all =
|
|
reinterpret_cast<const rocblas_double_complex*>(matrix_data);
|
|
const rocblas_double_complex* B_row_all =
|
|
reinterpret_cast<const rocblas_double_complex*>(rhs_data);
|
|
rocblas_double_complex* X_row_all =
|
|
reinterpret_cast<rocblas_double_complex*>(out_data);
|
|
|
|
auto dA_col_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(), A_batch_bytes, stream);
|
|
auto dB_col_alloc =
|
|
memory_utils::Alloc(dev_ctx.GetPlace(), B_batch_bytes, stream);
|
|
rocblas_double_complex* dA_col =
|
|
reinterpret_cast<rocblas_double_complex*>(dA_col_alloc->ptr());
|
|
rocblas_double_complex* dB_col =
|
|
reinterpret_cast<rocblas_double_complex*>(dB_col_alloc->ptr());
|
|
|
|
auto d_pivots_alloc = memory_utils::Alloc(
|
|
dev_ctx.GetPlace(),
|
|
static_cast<size_t>(batch_count) * order * sizeof(rocblas_int),
|
|
stream);
|
|
rocblas_int* d_pivots = reinterpret_cast<rocblas_int*>(d_pivots_alloc->ptr());
|
|
|
|
auto d_info_alloc = memory_utils::Alloc(
|
|
dev_ctx.GetPlace(),
|
|
static_cast<size_t>(batch_count) * sizeof(rocblas_int),
|
|
stream);
|
|
rocblas_int* d_info = reinterpret_cast<rocblas_int*>(d_info_alloc->ptr());
|
|
|
|
// A_row layout: row-major (order x order), B_row layout: row-major (order x
|
|
// rhs_cols)
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
const rocblas_double_complex* A_row =
|
|
A_row_all + static_cast<size_t>(i) * order * order;
|
|
rocblas_double_complex* A_col =
|
|
dA_col + static_cast<size_t>(i) * order * order;
|
|
const rocblas_double_complex* B_row =
|
|
B_row_all + static_cast<size_t>(i) * order * rhs_cols;
|
|
rocblas_double_complex* B_col =
|
|
dB_col + static_cast<size_t>(i) * order * rhs_cols;
|
|
|
|
// transpose A_row (row-major) -> A_col (column-major) via C = A^T
|
|
PADDLE_ENFORCE_GPU_SUCCESS(rocblas_zgeam(
|
|
rocblas_handle,
|
|
rocblas_operation_transpose,
|
|
rocblas_operation_none,
|
|
order,
|
|
order,
|
|
&kAlpha,
|
|
A_row,
|
|
order, // lda: when interpreting A_row as (order x order) row-major
|
|
&kZero,
|
|
nullptr,
|
|
order,
|
|
A_col,
|
|
order)); // ldc = order (column-major leading dim)
|
|
|
|
// transpose B_row (row-major order x rhs_cols) -> B_col (column-major order
|
|
// x rhs_cols)
|
|
PADDLE_ENFORCE_GPU_SUCCESS(rocblas_zgeam(
|
|
rocblas_handle,
|
|
rocblas_operation_transpose,
|
|
rocblas_operation_none,
|
|
order,
|
|
rhs_cols,
|
|
&kAlpha,
|
|
B_row,
|
|
rhs_cols, // lda when A_row is viewed row-major: leading = rhs_cols
|
|
&kZero,
|
|
nullptr,
|
|
rhs_cols,
|
|
B_col,
|
|
order)); // ldc = order
|
|
}
|
|
|
|
// LU factorization and solve for each batch
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
rocblas_double_complex* A_col =
|
|
dA_col + static_cast<size_t>(i) * order * order;
|
|
rocblas_double_complex* B_col =
|
|
dB_col + static_cast<size_t>(i) * order * rhs_cols;
|
|
rocblas_int* pivots_i = d_pivots + static_cast<size_t>(i) * order;
|
|
rocblas_int* info_i = d_info + i;
|
|
|
|
// getrf (LU factorization) on A_col (column-major)
|
|
PADDLE_ENFORCE_GPU_SUCCESS(rocsolver_zgetrf(
|
|
rocblas_handle, order, order, A_col, order, pivots_i, info_i));
|
|
|
|
// getrs: solve A_col * X_col = B_col
|
|
PADDLE_ENFORCE_GPU_SUCCESS(rocsolver_zgetrs(
|
|
rocblas_handle,
|
|
rocblas_operation_none, // no transpose on column-major matrix
|
|
order,
|
|
rhs_cols,
|
|
A_col,
|
|
order,
|
|
pivots_i,
|
|
B_col,
|
|
order));
|
|
}
|
|
|
|
// Transpose results back to row-major
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
rocblas_double_complex* B_col = dB_col + static_cast<size_t>(i) * order *
|
|
rhs_cols; // X in column-major
|
|
rocblas_double_complex* X_row =
|
|
X_row_all +
|
|
static_cast<size_t>(i) * order * rhs_cols; // target row-major
|
|
|
|
// transpose X_col -> X_row
|
|
PADDLE_ENFORCE_GPU_SUCCESS(rocblas_zgeam(
|
|
rocblas_handle,
|
|
rocblas_operation_transpose,
|
|
rocblas_operation_none,
|
|
rhs_cols,
|
|
order, // rowsC = rhs_cols, colsC = order
|
|
&kAlpha,
|
|
B_col,
|
|
order, // B_col lda = order (col-major)
|
|
&kZero,
|
|
nullptr,
|
|
order,
|
|
X_row,
|
|
rhs_cols)); // X_row ldc = rhs_cols (row-major leading dimension)
|
|
}
|
|
CPUPlace cpu_place;
|
|
DeviceContextPool& pool = DeviceContextPool::Instance();
|
|
auto* cpu_ctx = static_cast<CPUContext*>(pool.Get(cpu_place));
|
|
|
|
std::vector<rocblas_int> h_info(batch_count, 0);
|
|
memory_utils::Copy(CPUPlace(),
|
|
h_info.data(),
|
|
dev_ctx.GetPlace(),
|
|
d_info,
|
|
static_cast<size_t>(batch_count) * sizeof(rocblas_int),
|
|
reinterpret_cast<void*>(dev_ctx.stream()));
|
|
dev_ctx.Wait();
|
|
|
|
for (int i = 0; i < batch_count; ++i) {
|
|
PADDLE_ENFORCE_EQ(
|
|
h_info[i],
|
|
0,
|
|
errors::External("rocSOLVER getrf/getrs failed at batch %d, info: %d",
|
|
i,
|
|
h_info[i]));
|
|
}
|
|
}
|
|
#endif // PADDLE_WITH_HIP
|
|
|
|
template <typename T, typename Context>
|
|
void ComputeBackwardForComplexInputGPU(const DenseTensor& L,
|
|
const DenseTensor& V,
|
|
const optional<DenseTensor>& gL,
|
|
const optional<DenseTensor>& gV,
|
|
T* x_grad_data,
|
|
int batch_count,
|
|
int order,
|
|
const Context& dev_ctx) {
|
|
DenseTensor gL_safe;
|
|
if (gL.get_ptr()) {
|
|
gL_safe = gL.get();
|
|
} else {
|
|
gL_safe = Fill<T, Context>(dev_ctx, vectorize<int64_t>(L.dims()), T(0));
|
|
}
|
|
|
|
DenseTensor gV_safe;
|
|
if (gV.get_ptr()) {
|
|
gV_safe = gV.get();
|
|
} else {
|
|
gV_safe = Fill<T, Context>(dev_ctx, vectorize<int64_t>(V.dims()), T(0));
|
|
}
|
|
DenseTensor trans_v = TransposeLast2Dim<T>(dev_ctx, V);
|
|
DenseTensor Vh = Conj<T>(dev_ctx, trans_v);
|
|
DenseTensor Lconj = Conj<T>(dev_ctx, L);
|
|
DenseTensor Econj = Subtract<T>(
|
|
dev_ctx, funcs::Unsqueeze(Lconj, -2), funcs::Unsqueeze(Lconj, -1));
|
|
DenseTensor VhgV = Matmul<T>(dev_ctx, Vh, gV_safe);
|
|
DenseTensor diag_real = Real<T>(dev_ctx, VhgV);
|
|
|
|
auto cpu_place = CPUPlace();
|
|
DeviceContextPool& pool = DeviceContextPool::Instance();
|
|
auto* cpu_ctx = static_cast<CPUContext*>(pool.Get(cpu_place));
|
|
|
|
DenseTensor diag_real_cpu;
|
|
diag_real_cpu.Resize(diag_real.dims());
|
|
Copy(dev_ctx, diag_real, cpu_place, false, &diag_real_cpu);
|
|
|
|
DenseTensor diag_res_cpu =
|
|
funcs::BatchDiag<T>((*cpu_ctx), diag_real_cpu, batch_count);
|
|
|
|
DenseTensor diag_res;
|
|
dev_ctx.template Alloc<T>(&diag_res);
|
|
Copy(dev_ctx, diag_res_cpu, GPUPlace(), false, &diag_res);
|
|
|
|
DenseTensor diag_unsqueezed = funcs::Unsqueeze(diag_res, -2);
|
|
|
|
auto numel = diag_unsqueezed.numel();
|
|
DenseTensor diag_unsqueezed_complex;
|
|
auto* data_diag_un = diag_unsqueezed.data<dtype::Real<T>>();
|
|
diag_unsqueezed_complex.Resize(diag_unsqueezed.dims());
|
|
auto* data_diag_un_com = dev_ctx.template Alloc<T>(
|
|
&diag_unsqueezed_complex, static_cast<size_t>(numel * sizeof(T)));
|
|
|
|
funcs::ForRange<Context> for_range(dev_ctx, numel);
|
|
funcs::RealToComplexFunctor<T> functor(data_diag_un, data_diag_un_com, numel);
|
|
for_range(functor);
|
|
// real tensor multiply complex tensor in broadcast manner
|
|
DenseTensor res1 = Multiply<T>(dev_ctx, V, diag_unsqueezed_complex);
|
|
DenseTensor res2 = Matmul<T>(dev_ctx, Vh, res1);
|
|
DenseTensor result = Subtract<T>(dev_ctx, VhgV, res2);
|
|
|
|
result.Resize(V.dims());
|
|
dev_ctx.template Alloc<T>(&result);
|
|
result = Divide<T>(dev_ctx, result, Econj);
|
|
result =
|
|
funcs::DiagFill<T, T>(dev_ctx, order, order, order, 0, gL_safe, result);
|
|
DenseTensor rhs = Matmul<T>(dev_ctx, result, Vh);
|
|
|
|
// solve linear system
|
|
// solve(Vh, rhs, out, m, k)
|
|
// Vh: matrix with shape [m,m]
|
|
// rhs: rhs with shape [m,k]
|
|
// x_grad: out
|
|
int64_t m = Vh.dims(-1);
|
|
int64_t k = rhs.dims(-1);
|
|
auto* matrix_data = Vh.data<T>();
|
|
auto* rhs_data = rhs.data<T>();
|
|
|
|
SolveLinearSystemGPU<T>(
|
|
dev_ctx, matrix_data, rhs_data, x_grad_data, m, k, batch_count);
|
|
}
|
|
#endif // PADDLE_WITH_MAGMA
|
|
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
template <typename T, typename Context>
|
|
void EigGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& out_w,
|
|
const DenseTensor& out_v,
|
|
const optional<DenseTensor>& dout_w,
|
|
const optional<DenseTensor>& dout_v,
|
|
DenseTensor* dx) {
|
|
auto* dx_data = dev_ctx.template Alloc<dtype::Complex<T>>(dx);
|
|
if (dx->numel() == 0) {
|
|
return;
|
|
}
|
|
auto& dims = out_v.dims();
|
|
int batch_count = BatchCount(out_v);
|
|
const int64_t order = out_v.dims(-1);
|
|
|
|
ComputeBackwardForComplexInputGPU<dtype::Complex<T>, Context>(
|
|
out_w, out_v, dout_w, dout_v, dx_data, batch_count, order, dev_ctx);
|
|
}
|
|
#endif // PADDLE_WITH_CUDA || PADDLE_WITH_HIP
|
|
|
|
} // namespace phi
|
|
|
|
// Register the kernel
|
|
#ifdef PADDLE_WITH_MAGMA
|
|
PD_REGISTER_KERNEL(eig_grad,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::EigGradKernel,
|
|
float,
|
|
double,
|
|
phi::complex64,
|
|
phi::complex128) {
|
|
kernel->InputAt(0).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
|
|
kernel->InputAt(2).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
|
|
kernel->OutputAt(0).SetDataType(phi::dtype::ToComplex(kernel_key.dtype()));
|
|
}
|
|
#endif // PADDLE_WITH_MAGMA
|