chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,956 @@
|
||||
// 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
|
||||
Reference in New Issue
Block a user