wehub-resource-sync
957 lines
38 KiB
Plaintext
957 lines
38 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 <thrust/device_vector.h>
|
|
#include <algorithm>
|
|
#include <vector>
|
|
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/infermeta/unary.h"
|
|
#include "paddle/phi/kernels/diagonal_kernel.h"
|
|
#include "paddle/phi/kernels/fill_diagonal_tensor_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/complex_functors.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
#include "paddle/phi/kernels/funcs/parse_qr_mode.h"
|
|
#include "paddle/phi/kernels/impl/qr_kernel_impl.h"
|
|
#include "paddle/phi/kernels/qr_kernel.h"
|
|
#include "paddle/phi/kernels/slice_kernel.h"
|
|
#include "paddle/phi/kernels/transpose_kernel.h"
|
|
#include "paddle/phi/kernels/tril_triu_kernel.h"
|
|
|
|
namespace phi {
|
|
|
|
template <class T, class Context>
|
|
static DenseTensor Fill(const Context& dev_ctx,
|
|
std::vector<int64_t> shape,
|
|
T fill_value) {
|
|
DenseTensor ret;
|
|
ret.Resize(shape);
|
|
dev_ctx.template Alloc<T>(&ret);
|
|
funcs::SetConstant<Context, T>()(dev_ctx, &ret, fill_value);
|
|
return ret;
|
|
}
|
|
|
|
template <class T, class Context>
|
|
static DenseTensor identity_matrix(const Context& dev_ctx, common::DDim shape) {
|
|
DenseTensor M = Fill<T, Context>(dev_ctx, vectorize<int64_t>(shape), T(0));
|
|
size_t rank = M.dims().size();
|
|
int64_t M_diag_len = std::min(M.dims()[rank - 1], M.dims()[rank - 2]);
|
|
std::vector<int64_t> M_diag_shape;
|
|
for (size_t i = 0; i < rank - 2; ++i) {
|
|
M_diag_shape.push_back(M.dims()[i]);
|
|
}
|
|
M_diag_shape.push_back(M_diag_len);
|
|
DenseTensor M_diag = Fill<T, Context>(
|
|
dev_ctx, vectorize<int64_t>(make_ddim(M_diag_shape)), T(1));
|
|
M = FillDiagonalTensor<T, Context>(dev_ctx, M, M_diag, 0, rank - 2, rank - 1);
|
|
return M;
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
struct QrFunctor {
|
|
void operator()(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
bool compute_q,
|
|
bool reduced_mode,
|
|
DenseTensor* q,
|
|
DenseTensor* r) {
|
|
auto x_dims = x.dims();
|
|
int x_rank = x_dims.size();
|
|
int m = x_dims[x_rank - 2];
|
|
int n = x_dims[x_rank - 1];
|
|
int min_mn = std::min(m, n);
|
|
int k = reduced_mode ? min_mn : m;
|
|
int64_t batch_size =
|
|
static_cast<int64_t>(x.numel() / (static_cast<int64_t>(m) * n));
|
|
int64_t qr_stride = static_cast<int64_t>(m) * n;
|
|
int tau_stride = min_mn;
|
|
|
|
if (compute_q) {
|
|
dev_ctx.template Alloc<dtype::Real<T>>(
|
|
q, batch_size * m * k * sizeof(dtype::Real<T>));
|
|
}
|
|
dev_ctx.template Alloc<dtype::Real<T>>(
|
|
r, batch_size * k * n * sizeof(dtype::Real<T>));
|
|
|
|
// Note: allocate temporary tensors because of lacking in-place operations.
|
|
// Prepare qr
|
|
DenseTensor qr;
|
|
dev_ctx.template Alloc<dtype::Real<T>>(
|
|
&qr, size_t(batch_size * m * n * sizeof(dtype::Real<T>)));
|
|
// BatchedGeqrf performs computation in-place and 'qr' must be a copy of
|
|
// input
|
|
Copy(dev_ctx, x, dev_ctx.GetPlace(), false, &qr);
|
|
|
|
// Prepare tau
|
|
auto tau_dims_vec = vectorize<int64_t>(x_dims);
|
|
tau_dims_vec.pop_back();
|
|
tau_dims_vec[tau_dims_vec.size() - 1] = min_mn;
|
|
DenseTensor tau = Fill<T, Context>(dev_ctx, tau_dims_vec, T(0));
|
|
|
|
// Transpose 'qr' to conform the column-major order
|
|
auto tmp_qr = TransposeLast2Dim<T, Context>(dev_ctx, qr);
|
|
Copy(dev_ctx, tmp_qr, qr.place(), false, &qr);
|
|
auto qr_data = dev_ctx.template Alloc<dtype::Real<T>>(&qr);
|
|
auto tau_data = dev_ctx.template Alloc<dtype::Real<T>>(&tau);
|
|
|
|
BatchedGeqrf<Context, T>(
|
|
dev_ctx, batch_size, m, n, qr_data, m, tau_data, qr_stride, tau_stride);
|
|
|
|
if (reduced_mode) {
|
|
auto trans_qr = TransposeLast2Dim<T, Context>(dev_ctx, qr);
|
|
auto sliced_qr = Slice<T, Context>(
|
|
dev_ctx, trans_qr, {trans_qr.dims().size() - 2}, {0}, {min_mn});
|
|
auto tmp_r = TrilTriu<T, Context>(dev_ctx, sliced_qr, 0, false);
|
|
// Transpose 'tmp_r' to restore the original row-major order
|
|
Copy(dev_ctx, tmp_r, r->place(), false, r);
|
|
} else {
|
|
auto trans_qr = TransposeLast2Dim<T, Context>(dev_ctx, qr);
|
|
auto tmp_r = TrilTriu<T, Context>(dev_ctx, trans_qr, 0, false);
|
|
// Transpose 'tmp_r' to restore the original row-major order
|
|
Copy(dev_ctx, tmp_r, r->place(), false, r);
|
|
}
|
|
|
|
if (compute_q) {
|
|
// Perform QRGQR for Q using the result from GEQRF
|
|
// Transpose 'q' to restore the original row-major order
|
|
if (reduced_mode) {
|
|
BatchedOrgqr<Context, T>(dev_ctx,
|
|
batch_size,
|
|
m,
|
|
min_mn,
|
|
min_mn,
|
|
qr_data,
|
|
m,
|
|
tau_data,
|
|
qr_stride,
|
|
tau_stride);
|
|
auto trans_q = TransposeLast2Dim<T, Context>(dev_ctx, qr);
|
|
auto sliced_q = Slice<T, Context>(
|
|
dev_ctx, trans_q, {trans_q.dims().size() - 1}, {0}, {min_mn});
|
|
Copy(dev_ctx, sliced_q, q->place(), false, q);
|
|
} else {
|
|
if (m > n) {
|
|
auto new_qr_dims_vec = vectorize<int64_t>(x_dims);
|
|
new_qr_dims_vec[new_qr_dims_vec.size() - 1] = m;
|
|
DenseTensor new_qr = Fill<T, Context>(dev_ctx, new_qr_dims_vec, T(0));
|
|
auto new_qr_data = dev_ctx.template Alloc<dtype::Real<T>>(&new_qr);
|
|
auto new_qr_stride = m * m;
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
memory_utils::Copy(dev_ctx.GetPlace(),
|
|
(new_qr_data + i * new_qr_stride),
|
|
dev_ctx.GetPlace(),
|
|
(qr_data + i * qr_stride),
|
|
qr_stride * sizeof(dtype::Real<T>),
|
|
dev_ctx.stream());
|
|
}
|
|
BatchedOrgqr<Context, T>(dev_ctx,
|
|
batch_size,
|
|
m,
|
|
m,
|
|
min_mn,
|
|
new_qr_data,
|
|
m,
|
|
tau_data,
|
|
new_qr_stride,
|
|
tau_stride);
|
|
auto trans_q = TransposeLast2Dim<T, Context>(dev_ctx, new_qr);
|
|
Copy(dev_ctx, trans_q, q->place(), false, q);
|
|
} else {
|
|
BatchedOrgqr<Context, T>(dev_ctx,
|
|
batch_size,
|
|
m,
|
|
m,
|
|
min_mn,
|
|
qr_data,
|
|
m,
|
|
tau_data,
|
|
qr_stride,
|
|
tau_stride);
|
|
auto trans_q = TransposeLast2Dim<T, Context>(dev_ctx, qr);
|
|
auto sliced_q = Slice<T, Context>(
|
|
dev_ctx, trans_q, {trans_q.dims().size() - 1}, {0}, {m});
|
|
Copy(dev_ctx, sliced_q, q->place(), false, q);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename T, typename Context>
|
|
struct QrFunctor<dtype::complex<T>, Context> {
|
|
void operator()(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
bool compute_q,
|
|
bool reduced_mode,
|
|
DenseTensor* q,
|
|
DenseTensor* r) {
|
|
auto x_dims = x.dims();
|
|
int x_rank = x_dims.size();
|
|
int m = x_dims[x_rank - 2];
|
|
int n = x_dims[x_rank - 1];
|
|
int min_mn = std::min(m, n);
|
|
int k = reduced_mode ? min_mn : m;
|
|
int64_t batch_size = x.numel() / (static_cast<int64_t>(m) * n);
|
|
int64_t qr_stride = static_cast<int64_t>(m) * n;
|
|
int tau_stride = min_mn;
|
|
if (compute_q) {
|
|
dev_ctx.template Alloc<dtype::complex<T>>(
|
|
q, batch_size * m * k * sizeof(dtype::complex<T>));
|
|
}
|
|
dev_ctx.template Alloc<dtype::complex<T>>(
|
|
r, batch_size * k * n * sizeof(dtype::complex<T>));
|
|
// Note: allocate temporary tensors because of lacking in-place operations.
|
|
// Prepare qr
|
|
DenseTensor qr;
|
|
dev_ctx.template Alloc<dtype::complex<T>>(
|
|
&qr, size_t(batch_size * m * n * sizeof(dtype::complex<T>)));
|
|
// BatchedGeqrf performs computation in-place and 'qr' must be a copy of
|
|
// input
|
|
Copy(dev_ctx, x, dev_ctx.GetPlace(), false, &qr);
|
|
// Prepare tau
|
|
auto tau_dims_vec = vectorize<int64_t>(x_dims);
|
|
tau_dims_vec.pop_back();
|
|
tau_dims_vec[tau_dims_vec.size() - 1] = min_mn;
|
|
DenseTensor tau =
|
|
Fill<dtype::complex<T>, Context>(dev_ctx, tau_dims_vec, T(0));
|
|
// Transpose 'qr' to conform the column-major order
|
|
auto tmp_qr = TransposeLast2Dim<dtype::complex<T>, Context>(dev_ctx, qr);
|
|
Copy(dev_ctx, tmp_qr, qr.place(), false, &qr);
|
|
auto qr_data = dev_ctx.template Alloc<dtype::complex<T>>(&qr);
|
|
auto tau_data = dev_ctx.template Alloc<dtype::complex<T>>(&tau);
|
|
BatchedGeqrf<Context, dtype::complex<T>>(
|
|
dev_ctx, batch_size, m, n, qr_data, m, tau_data, qr_stride, tau_stride);
|
|
if (reduced_mode) {
|
|
auto trans_qr =
|
|
TransposeLast2Dim<dtype::complex<T>, Context>(dev_ctx, qr);
|
|
auto sliced_qr = Slice<dtype::complex<T>, Context>(
|
|
dev_ctx, trans_qr, {trans_qr.dims().size() - 2}, {0}, {min_mn});
|
|
auto tmp_r =
|
|
TrilTriu<dtype::complex<T>, Context>(dev_ctx, sliced_qr, 0, false);
|
|
// Transpose 'tmp_r' to restore the original row-major order
|
|
Copy(dev_ctx, tmp_r, r->place(), false, r);
|
|
} else {
|
|
auto trans_qr =
|
|
TransposeLast2Dim<dtype::complex<T>, Context>(dev_ctx, qr);
|
|
auto tmp_r =
|
|
TrilTriu<dtype::complex<T>, Context>(dev_ctx, trans_qr, 0, false);
|
|
// Transpose 'tmp_r' to restore the original row-major order
|
|
Copy(dev_ctx, tmp_r, r->place(), false, r);
|
|
}
|
|
if (compute_q) {
|
|
// Perform QRGQR for Q using the result from GEQRF
|
|
// Transpose 'q' to restore the original row-major order
|
|
if (reduced_mode) {
|
|
BatchedOrgqr<Context, dtype::complex<T>>(dev_ctx,
|
|
batch_size,
|
|
m,
|
|
min_mn,
|
|
min_mn,
|
|
qr_data,
|
|
m,
|
|
tau_data,
|
|
qr_stride,
|
|
tau_stride);
|
|
auto trans_q =
|
|
TransposeLast2Dim<dtype::complex<T>, Context>(dev_ctx, qr);
|
|
auto sliced_q = Slice<dtype::complex<T>, Context>(
|
|
dev_ctx, trans_q, {trans_q.dims().size() - 1}, {0}, {min_mn});
|
|
Copy(dev_ctx, sliced_q, q->place(), false, q);
|
|
} else {
|
|
if (m > n) {
|
|
auto new_qr_dims_vec = vectorize<int64_t>(x_dims);
|
|
new_qr_dims_vec[new_qr_dims_vec.size() - 1] = m;
|
|
DenseTensor new_qr =
|
|
Fill<dtype::complex<T>, Context>(dev_ctx, new_qr_dims_vec, T(0));
|
|
auto new_qr_data = dev_ctx.template Alloc<dtype::complex<T>>(&new_qr);
|
|
auto new_qr_stride = m * m;
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
memory_utils::Copy(dev_ctx.GetPlace(),
|
|
(new_qr_data + i * new_qr_stride),
|
|
dev_ctx.GetPlace(),
|
|
(qr_data + i * qr_stride),
|
|
qr_stride * sizeof(dtype::complex<T>),
|
|
dev_ctx.stream());
|
|
}
|
|
BatchedOrgqr<Context, dtype::complex<T>>(dev_ctx,
|
|
batch_size,
|
|
m,
|
|
m,
|
|
min_mn,
|
|
new_qr_data,
|
|
m,
|
|
tau_data,
|
|
new_qr_stride,
|
|
tau_stride);
|
|
auto trans_q =
|
|
TransposeLast2Dim<dtype::complex<T>, Context>(dev_ctx, new_qr);
|
|
Copy(dev_ctx, trans_q, q->place(), false, q);
|
|
} else {
|
|
BatchedOrgqr<Context, dtype::complex<T>>(dev_ctx,
|
|
batch_size,
|
|
m,
|
|
m,
|
|
min_mn,
|
|
qr_data,
|
|
m,
|
|
tau_data,
|
|
qr_stride,
|
|
tau_stride);
|
|
auto trans_q =
|
|
TransposeLast2Dim<dtype::complex<T>, Context>(dev_ctx, qr);
|
|
auto sliced_q = Slice<dtype::complex<T>, Context>(
|
|
dev_ctx, trans_q, {trans_q.dims().size() - 1}, {0}, {m});
|
|
Copy(dev_ctx, sliced_q, q->place(), false, q);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename T, typename Context>
|
|
void QrKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const std::string& mode,
|
|
DenseTensor* q,
|
|
DenseTensor* r) {
|
|
bool compute_q;
|
|
bool reduced_mode;
|
|
std::tie(compute_q, reduced_mode) = funcs::ParseQrMode(mode);
|
|
if (x.numel() == 0) {
|
|
if (q->numel() == 0) {
|
|
q->Resize(q->dims());
|
|
} else {
|
|
*q = identity_matrix<T, Context>(dev_ctx, q->dims());
|
|
}
|
|
r->Resize(r->dims());
|
|
dev_ctx.template Alloc<T>(q);
|
|
dev_ctx.template Alloc<T>(r);
|
|
return;
|
|
}
|
|
QrFunctor<T, Context>()(dev_ctx, x, compute_q, reduced_mode, q, r);
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
#define FUNC_WITH_TYPES(m) m(float, s) m(double, d)
|
|
#define GEQRF_BATCH_INSTANCE(T, C) \
|
|
template <> \
|
|
void BatchedGeqrf<GPUContext, T>(const GPUContext& dev_ctx, \
|
|
int batch_size, \
|
|
int m, \
|
|
int n, \
|
|
T* a, \
|
|
int lda, \
|
|
T* tau, \
|
|
int a_stride, \
|
|
int tau_stride) { \
|
|
auto handle = dev_ctx.cusolver_dn_handle(); \
|
|
for (int i = 0; i < batch_size; ++i) { \
|
|
T* a_working_ptr = &a[i * a_stride]; \
|
|
T* tau_working_ptr = &tau[i * tau_stride]; \
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::rocsolver_##C##geqrf( \
|
|
handle, m, n, a_working_ptr, lda, tau_working_ptr)); \
|
|
} \
|
|
}
|
|
|
|
FUNC_WITH_TYPES(GEQRF_BATCH_INSTANCE);
|
|
|
|
#define ORGQR_BATCH_INSTANCE(T, C) \
|
|
template <> \
|
|
void BatchedOrgqr<GPUContext, T>(const GPUContext& dev_ctx, \
|
|
int batch_size, \
|
|
int m, \
|
|
int n, \
|
|
int k, \
|
|
T* a, \
|
|
int lda, \
|
|
T* tau, \
|
|
int a_stride, \
|
|
int tau_stride) { \
|
|
auto handle = dev_ctx.cusolver_dn_handle(); \
|
|
for (int i = 0; i < batch_size; ++i) { \
|
|
T* a_working_ptr = &a[i * a_stride]; \
|
|
T* tau_working_ptr = &tau[i * tau_stride]; \
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::rocsolver_##C##orgqr( \
|
|
handle, m, n, k, a_working_ptr, lda, tau_working_ptr)); \
|
|
} \
|
|
}
|
|
|
|
FUNC_WITH_TYPES(ORGQR_BATCH_INSTANCE);
|
|
#else
|
|
template <>
|
|
void BatchedGeqrf<GPUContext, float>(const GPUContext& dev_ctx,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
float* a,
|
|
int lda,
|
|
float* tau,
|
|
int a_stride,
|
|
int tau_stride) {
|
|
if (static_cast<int64_t>(m) * n * 171 > std::numeric_limits<int>::max()) {
|
|
const int64_t batch_size_64 = static_cast<int64_t>(batch_size);
|
|
const int64_t m_64 = static_cast<int64_t>(m);
|
|
const int64_t n_64 = static_cast<int64_t>(n);
|
|
const int64_t lda_64 = static_cast<int64_t>(lda);
|
|
const int64_t a_stride_64 = static_cast<int64_t>(a_stride);
|
|
const int64_t tau_stride_64 = static_cast<int64_t>(tau_stride);
|
|
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
|
|
size_t workspace_in_bytes_on_device = 0;
|
|
size_t workspace_in_bytes_on_host = 0;
|
|
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnXgeqrf_bufferSize(handle,
|
|
nullptr,
|
|
m_64,
|
|
n_64,
|
|
CUDA_R_32F,
|
|
a,
|
|
lda_64,
|
|
CUDA_R_32F,
|
|
tau,
|
|
CUDA_R_32F,
|
|
&workspace_in_bytes_on_device,
|
|
&workspace_in_bytes_on_host));
|
|
|
|
DenseTensor device_workspace;
|
|
device_workspace.Resize(
|
|
make_ddim({static_cast<int64_t>(workspace_in_bytes_on_device)}));
|
|
uint8_t* device_workspace_ptr =
|
|
dev_ctx.template Alloc<uint8_t>(&device_workspace);
|
|
|
|
DenseTensor host_workspace;
|
|
uint8_t* host_workspace_ptr = nullptr;
|
|
|
|
if (workspace_in_bytes_on_host > 0) {
|
|
host_workspace.Resize(
|
|
make_ddim({static_cast<int64_t>(workspace_in_bytes_on_host)}));
|
|
host_workspace_ptr = dev_ctx.template HostAlloc<uint8_t>(&host_workspace);
|
|
}
|
|
|
|
DenseTensor info;
|
|
info.Resize({1});
|
|
int* info_d = dev_ctx.template Alloc<int>(&info);
|
|
|
|
for (int64_t i = 0; i < batch_size_64; ++i) {
|
|
float* a_working_ptr = &a[i * a_stride_64];
|
|
float* tau_working_ptr = &tau[i * tau_stride_64];
|
|
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnXgeqrf(handle,
|
|
nullptr,
|
|
m_64,
|
|
n_64,
|
|
CUDA_R_32F,
|
|
a_working_ptr,
|
|
lda_64,
|
|
CUDA_R_32F,
|
|
tau_working_ptr,
|
|
CUDA_R_32F,
|
|
device_workspace_ptr,
|
|
workspace_in_bytes_on_device,
|
|
host_workspace_ptr,
|
|
workspace_in_bytes_on_host,
|
|
info_d));
|
|
|
|
int info_h;
|
|
memory_utils::Copy(CPUPlace(),
|
|
&info_h,
|
|
dev_ctx.GetPlace(),
|
|
info_d,
|
|
sizeof(int),
|
|
dev_ctx.stream());
|
|
PADDLE_ENFORCE_EQ(
|
|
info_h,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"For batch [%d]: CUSolver (64-bit) geqrf is not zero. [%d]",
|
|
i,
|
|
info_h));
|
|
}
|
|
} else {
|
|
int lwork = 0;
|
|
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnSgeqrf_bufferSize(handle, m, n, a, lda, &lwork));
|
|
|
|
DenseTensor workspace = DenseTensor();
|
|
workspace.Resize({lwork});
|
|
float* workspace_ptr = dev_ctx.template Alloc<float>(&workspace);
|
|
|
|
DenseTensor info = DenseTensor();
|
|
info.Resize({1});
|
|
int* info_d = dev_ctx.template Alloc<int>(&info);
|
|
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
float* a_working_ptr = &a[i * a_stride];
|
|
float* tau_working_ptr = &tau[i * tau_stride];
|
|
// compute geqrf
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnSgeqrf(handle,
|
|
m,
|
|
n,
|
|
a_working_ptr,
|
|
lda,
|
|
tau_working_ptr,
|
|
workspace_ptr,
|
|
lwork,
|
|
info_d));
|
|
// Do we need synchronized here?
|
|
// check the error info
|
|
int info_h;
|
|
memory_utils::Copy(CPUPlace(),
|
|
&info_h,
|
|
dev_ctx.GetPlace(),
|
|
info_d,
|
|
sizeof(int),
|
|
dev_ctx.stream());
|
|
PADDLE_ENFORCE_EQ(
|
|
info_h,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"For batch [%d]: CUSolver geqrf is not zero. [%d]", i, info_h));
|
|
}
|
|
}
|
|
}
|
|
|
|
template <>
|
|
void BatchedGeqrf<GPUContext, double>(const GPUContext& dev_ctx,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
double* a,
|
|
int lda,
|
|
double* tau,
|
|
int a_stride,
|
|
int tau_stride) {
|
|
int lwork = 0;
|
|
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnDgeqrf_bufferSize(handle, m, n, a, lda, &lwork));
|
|
|
|
DenseTensor workspace = DenseTensor();
|
|
workspace.Resize({lwork});
|
|
double* workspace_ptr = dev_ctx.template Alloc<double>(&workspace);
|
|
|
|
DenseTensor info = DenseTensor();
|
|
info.Resize({1});
|
|
int* info_d = dev_ctx.template Alloc<int>(&info);
|
|
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
double* a_working_ptr = &a[i * a_stride];
|
|
double* tau_working_ptr = &tau[i * tau_stride];
|
|
// compute geqrf
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnDgeqrf(handle,
|
|
m,
|
|
n,
|
|
a_working_ptr,
|
|
lda,
|
|
tau_working_ptr,
|
|
workspace_ptr,
|
|
lwork,
|
|
info_d));
|
|
// Do we need synchronized here?
|
|
// check the error info
|
|
int info_h;
|
|
memory_utils::Copy(CPUPlace(),
|
|
&info_h,
|
|
dev_ctx.GetPlace(),
|
|
info_d,
|
|
sizeof(int),
|
|
dev_ctx.stream());
|
|
PADDLE_ENFORCE_EQ(
|
|
info_h,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"For batch [%d]: CUSolver geqrf is not zero. [%d]", i, info_h));
|
|
}
|
|
}
|
|
|
|
template <>
|
|
void BatchedGeqrf<GPUContext, complex64>(const GPUContext& dev_ctx,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
complex64* a,
|
|
int lda,
|
|
complex64* tau,
|
|
int a_stride,
|
|
int tau_stride) {
|
|
int lwork = 0;
|
|
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnCgeqrf_bufferSize(
|
|
handle, m, n, reinterpret_cast<cuComplex*>(a), lda, &lwork));
|
|
|
|
DenseTensor workspace = DenseTensor();
|
|
workspace.Resize({lwork});
|
|
complex64* workspace_ptr = dev_ctx.template Alloc<complex64>(&workspace);
|
|
|
|
DenseTensor info = DenseTensor();
|
|
info.Resize({1});
|
|
int* info_d = dev_ctx.template Alloc<int>(&info);
|
|
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
complex64* a_working_ptr = &a[i * a_stride];
|
|
complex64* tau_working_ptr = &tau[i * tau_stride];
|
|
// compute geqrf
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnCgeqrf(handle,
|
|
m,
|
|
n,
|
|
reinterpret_cast<cuComplex*>(a_working_ptr),
|
|
lda,
|
|
reinterpret_cast<cuComplex*>(tau_working_ptr),
|
|
reinterpret_cast<cuComplex*>(workspace_ptr),
|
|
lwork,
|
|
info_d));
|
|
// Do we need synchronized here?
|
|
// check the error info
|
|
int info_h;
|
|
memory_utils::Copy(CPUPlace(),
|
|
&info_h,
|
|
dev_ctx.GetPlace(),
|
|
info_d,
|
|
sizeof(int),
|
|
dev_ctx.stream());
|
|
PADDLE_ENFORCE_EQ(
|
|
info_h,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"For batch [%d]: CUSolver geqrf is not zero. [%d]", i, info_h));
|
|
}
|
|
}
|
|
|
|
template <>
|
|
void BatchedGeqrf<GPUContext, complex128>(const GPUContext& dev_ctx,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
complex128* a,
|
|
int lda,
|
|
complex128* tau,
|
|
int a_stride,
|
|
int tau_stride) {
|
|
int lwork = 0;
|
|
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnZgeqrf_bufferSize(
|
|
handle, m, n, reinterpret_cast<cuDoubleComplex*>(a), lda, &lwork));
|
|
|
|
DenseTensor workspace = DenseTensor();
|
|
workspace.Resize({lwork});
|
|
complex128* workspace_ptr = dev_ctx.template Alloc<complex128>(&workspace);
|
|
|
|
DenseTensor info = DenseTensor();
|
|
info.Resize({1});
|
|
int* info_d = dev_ctx.template Alloc<int>(&info);
|
|
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
complex128* a_working_ptr = &a[i * a_stride];
|
|
complex128* tau_working_ptr = &tau[i * tau_stride];
|
|
// compute geqrf
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnZgeqrf(
|
|
handle,
|
|
m,
|
|
n,
|
|
reinterpret_cast<cuDoubleComplex*>(a_working_ptr),
|
|
lda,
|
|
reinterpret_cast<cuDoubleComplex*>(tau_working_ptr),
|
|
reinterpret_cast<cuDoubleComplex*>(workspace_ptr),
|
|
lwork,
|
|
info_d));
|
|
// Do we need synchronized here?
|
|
// check the error info
|
|
int info_h;
|
|
memory_utils::Copy(CPUPlace(),
|
|
&info_h,
|
|
dev_ctx.GetPlace(),
|
|
info_d,
|
|
sizeof(int),
|
|
dev_ctx.stream());
|
|
PADDLE_ENFORCE_EQ(
|
|
info_h,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"For batch [%d]: CUSolver geqrf is not zero. [%d]", i, info_h));
|
|
}
|
|
}
|
|
|
|
template <>
|
|
void BatchedOrgqr<GPUContext, float>(const GPUContext& dev_ctx,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
float* a,
|
|
int lda,
|
|
float* tau,
|
|
int a_stride,
|
|
int tau_stride) {
|
|
int lwork = 0;
|
|
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnSorgqr_bufferSize(
|
|
handle, m, n, k, a, lda, tau, &lwork));
|
|
|
|
DenseTensor workspace = DenseTensor();
|
|
workspace.Resize({lwork});
|
|
float* workspace_ptr = dev_ctx.template Alloc<float>(&workspace);
|
|
|
|
DenseTensor info = DenseTensor();
|
|
info.Resize({1});
|
|
int* info_d = dev_ctx.template Alloc<int>(&info);
|
|
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
float* a_working_ptr = &a[i * a_stride];
|
|
float* tau_working_ptr = &tau[i * tau_stride];
|
|
// compute orggr
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnSorgqr(handle,
|
|
m,
|
|
n,
|
|
k,
|
|
a_working_ptr,
|
|
lda,
|
|
tau_working_ptr,
|
|
workspace_ptr,
|
|
lwork,
|
|
info_d));
|
|
// Do we need synchronized here?
|
|
// check the error info
|
|
int info_h;
|
|
memory_utils::Copy(CPUPlace(),
|
|
&info_h,
|
|
dev_ctx.GetPlace(),
|
|
info_d,
|
|
sizeof(int),
|
|
dev_ctx.stream());
|
|
PADDLE_ENFORCE_EQ(
|
|
info_h,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"For batch [%d]: CUSolver QR is not zero. [%d]", i, info_h));
|
|
}
|
|
}
|
|
|
|
template <>
|
|
void BatchedOrgqr<GPUContext, double>(const GPUContext& dev_ctx,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
double* a,
|
|
int lda,
|
|
double* tau,
|
|
int a_stride,
|
|
int tau_stride) {
|
|
int lwork = 0;
|
|
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnDorgqr_bufferSize(
|
|
handle, m, n, k, a, lda, tau, &lwork));
|
|
|
|
DenseTensor workspace = DenseTensor();
|
|
workspace.Resize({lwork});
|
|
double* workspace_ptr = dev_ctx.template Alloc<double>(&workspace);
|
|
|
|
DenseTensor info = DenseTensor();
|
|
info.Resize({1});
|
|
int* info_d = dev_ctx.template Alloc<int>(&info);
|
|
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
double* a_working_ptr = &a[i * a_stride];
|
|
double* tau_working_ptr = &tau[i * tau_stride];
|
|
// compute orggr
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnDorgqr(handle,
|
|
m,
|
|
n,
|
|
k,
|
|
a_working_ptr,
|
|
lda,
|
|
tau_working_ptr,
|
|
workspace_ptr,
|
|
lwork,
|
|
info_d));
|
|
// Do we need synchronized here?
|
|
// check the error info
|
|
int info_h;
|
|
memory_utils::Copy(CPUPlace(),
|
|
&info_h,
|
|
dev_ctx.GetPlace(),
|
|
info_d,
|
|
sizeof(int),
|
|
dev_ctx.stream());
|
|
PADDLE_ENFORCE_EQ(
|
|
info_h,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"For batch [%d]: CUSolver QR is not zero. [%d]", i, info_h));
|
|
}
|
|
}
|
|
|
|
template <>
|
|
void BatchedOrgqr<GPUContext, complex64>(const GPUContext& dev_ctx,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
complex64* a,
|
|
int lda,
|
|
complex64* tau,
|
|
int a_stride,
|
|
int tau_stride) {
|
|
int lwork = 0;
|
|
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnCungqr_bufferSize(handle,
|
|
m,
|
|
n,
|
|
k,
|
|
reinterpret_cast<cuComplex*>(a),
|
|
lda,
|
|
reinterpret_cast<cuComplex*>(tau),
|
|
&lwork));
|
|
|
|
DenseTensor workspace = DenseTensor();
|
|
workspace.Resize({lwork});
|
|
complex64* workspace_ptr = dev_ctx.template Alloc<complex64>(&workspace);
|
|
|
|
DenseTensor info = DenseTensor();
|
|
info.Resize({1});
|
|
int* info_d = dev_ctx.template Alloc<int>(&info);
|
|
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
complex64* a_working_ptr = &a[i * a_stride];
|
|
complex64* tau_working_ptr = &tau[i * tau_stride];
|
|
// compute orggr
|
|
PADDLE_ENFORCE_GPU_SUCCESS(
|
|
dynload::cusolverDnCungqr(handle,
|
|
m,
|
|
n,
|
|
k,
|
|
reinterpret_cast<cuComplex*>(a_working_ptr),
|
|
lda,
|
|
reinterpret_cast<cuComplex*>(tau_working_ptr),
|
|
reinterpret_cast<cuComplex*>(workspace_ptr),
|
|
lwork,
|
|
info_d));
|
|
// Do we need synchronized here?
|
|
// check the error info
|
|
int info_h;
|
|
memory_utils::Copy(CPUPlace(),
|
|
&info_h,
|
|
dev_ctx.GetPlace(),
|
|
info_d,
|
|
sizeof(int),
|
|
dev_ctx.stream());
|
|
PADDLE_ENFORCE_EQ(
|
|
info_h,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"For batch [%d]: CUSolver QR is not zero. [%d]", i, info_h));
|
|
}
|
|
}
|
|
|
|
template <>
|
|
void BatchedOrgqr<GPUContext, complex128>(const GPUContext& dev_ctx,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
complex128* a,
|
|
int lda,
|
|
complex128* tau,
|
|
int a_stride,
|
|
int tau_stride) {
|
|
int lwork = 0;
|
|
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnZungqr_bufferSize(
|
|
handle,
|
|
m,
|
|
n,
|
|
k,
|
|
reinterpret_cast<cuDoubleComplex*>(a),
|
|
lda,
|
|
reinterpret_cast<cuDoubleComplex*>(tau),
|
|
&lwork));
|
|
|
|
DenseTensor workspace = DenseTensor();
|
|
workspace.Resize({lwork});
|
|
complex128* workspace_ptr = dev_ctx.template Alloc<complex128>(&workspace);
|
|
|
|
DenseTensor info = DenseTensor();
|
|
info.Resize({1});
|
|
int* info_d = dev_ctx.template Alloc<int>(&info);
|
|
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
complex128* a_working_ptr = &a[i * a_stride];
|
|
complex128* tau_working_ptr = &tau[i * tau_stride];
|
|
// compute orggr
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnZungqr(
|
|
handle,
|
|
m,
|
|
n,
|
|
k,
|
|
reinterpret_cast<cuDoubleComplex*>(a_working_ptr),
|
|
lda,
|
|
reinterpret_cast<cuDoubleComplex*>(tau_working_ptr),
|
|
reinterpret_cast<cuDoubleComplex*>(workspace_ptr),
|
|
lwork,
|
|
info_d));
|
|
// Do we need synchronized here?
|
|
// check the error info
|
|
int info_h;
|
|
memory_utils::Copy(CPUPlace(),
|
|
&info_h,
|
|
dev_ctx.GetPlace(),
|
|
info_d,
|
|
sizeof(int),
|
|
dev_ctx.stream());
|
|
PADDLE_ENFORCE_EQ(
|
|
info_h,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"For batch [%d]: CUSolver QR is not zero. [%d]", i, info_h));
|
|
}
|
|
}
|
|
#endif
|
|
|
|
} // namespace phi
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
PD_REGISTER_KERNEL(qr, GPU, ALL_LAYOUT, phi::QrKernel, float, double) {}
|
|
#else
|
|
PD_REGISTER_KERNEL(qr,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::QrKernel,
|
|
float,
|
|
double,
|
|
phi::complex64,
|
|
phi::complex128) {}
|
|
#endif
|