wehub-resource-sync
325 lines
14 KiB
C++
325 lines
14 KiB
C++
// 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.
|
|
|
|
#pragma once
|
|
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/utils/optional.h"
|
|
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/kernels/activation_kernel.h"
|
|
#include "paddle/phi/kernels/elementwise_subtract_kernel.h"
|
|
#include "paddle/phi/kernels/matmul_kernel.h"
|
|
#include "paddle/phi/kernels/reduce_sum_kernel.h"
|
|
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
#include "paddle/phi/backends/dynload/cusolver.h"
|
|
#endif
|
|
|
|
#if defined(PADDLE_WITH_HIP)
|
|
#include "paddle/phi/backends/dynload/rocsolver.h"
|
|
#endif
|
|
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#endif
|
|
|
|
namespace phi {
|
|
|
|
inline int GetBatchCount(const DDim& dims) {
|
|
int64_t count = 1;
|
|
int num_dims = dims.size();
|
|
for (int i = 0; i < num_dims - 2; ++i) {
|
|
count *= dims[i];
|
|
}
|
|
PADDLE_ENFORCE_LE_INT_MAX(count, "lstsq batch count");
|
|
return static_cast<int>(count);
|
|
}
|
|
|
|
inline int GetMatrixStride(const DDim& dims) {
|
|
int num_dims = dims.size();
|
|
int64_t stride = dims[num_dims - 1] * dims[num_dims - 2];
|
|
PADDLE_ENFORCE_LE_INT_MAX(stride, "lstsq matrix stride");
|
|
return static_cast<int>(stride);
|
|
}
|
|
|
|
inline bool IsComplexDtype(const DataType& type) {
|
|
return (type == DataType::COMPLEX64 || type == DataType::COMPLEX128);
|
|
}
|
|
|
|
template <typename Context, typename T>
|
|
inline void GetResidualsTensor(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& y,
|
|
const std::string& driver,
|
|
DenseTensor* solution,
|
|
DenseTensor* residuals,
|
|
DenseTensor* rank) {
|
|
auto x_dims = x.dims();
|
|
int dim_size = x_dims.size();
|
|
int64_t m = x_dims[dim_size - 2];
|
|
int64_t n = x_dims[dim_size - 1];
|
|
|
|
// Note(zrr1999): Although m and n are declared as int64_t, the rank tensor
|
|
// stores int values (see rank->data<int>() usage below), so effectively these
|
|
// dimensions are limited to int range in the current implementation.
|
|
if (m > n && driver != "gelsy") {
|
|
bool compute_residuals = true;
|
|
if ((driver == "gelss" || driver == "gelsd") && rank->numel() != 0) {
|
|
if (dim_size == 2) {
|
|
compute_residuals = rank->data<int>()[0] == n;
|
|
} else {
|
|
compute_residuals = std::all_of(rank->data<int>(),
|
|
rank->data<int>() + rank->numel(),
|
|
[n](int r) { return r == n; });
|
|
}
|
|
}
|
|
if (compute_residuals) {
|
|
DenseTensor matmul_tensor =
|
|
Matmul<T>(dev_ctx, x, *solution, false, false);
|
|
DenseTensor sub_tensor = Subtract<T>(dev_ctx, matmul_tensor, y);
|
|
DenseTensor pow_tensor;
|
|
pow_tensor.Resize(sub_tensor.dims());
|
|
dev_ctx.template Alloc<T>(&pow_tensor);
|
|
PowKernel<T>(dev_ctx, sub_tensor, Scalar(2), &pow_tensor);
|
|
|
|
auto sum_tensor = Sum<T>(
|
|
dev_ctx, pow_tensor, IntArray({-2}), pow_tensor.dtype(), false);
|
|
Copy<Context>(dev_ctx, sum_tensor, dev_ctx.GetPlace(), true, residuals);
|
|
return;
|
|
}
|
|
}
|
|
|
|
IntArray empty_shape({0});
|
|
DenseTensor empty_tensor = Empty<T, Context>(dev_ctx, empty_shape);
|
|
Copy<Context>(dev_ctx, empty_tensor, dev_ctx.GetPlace(), true, residuals);
|
|
}
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
template <typename Context, typename T>
|
|
inline void BatchedOrmqr(const Context& dev_ctx,
|
|
bool left,
|
|
bool transpose,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
T* a,
|
|
int a_stride,
|
|
T* tau,
|
|
int tau_stride,
|
|
T* other,
|
|
int other_stride);
|
|
|
|
#define FUNC_WITH_TYPES(m) m(float, s) m(double, d)
|
|
#define ORMQR_BATCH_INSTANCE(T, C) \
|
|
template <> \
|
|
inline void BatchedOrmqr<GPUContext, T>(const GPUContext& dev_ctx, \
|
|
bool left, \
|
|
bool transpose, \
|
|
int batch_size, \
|
|
int m, \
|
|
int n, \
|
|
int k, \
|
|
T* a, \
|
|
int a_stride, \
|
|
T* tau, \
|
|
int tau_stride, \
|
|
T* other, \
|
|
int other_stride) { \
|
|
auto side = left ? rocblas_side_left : rocblas_side_right; \
|
|
auto trans = \
|
|
transpose ? rocblas_operation_transpose : rocblas_operation_none; \
|
|
int lda = std::max<int>(1, left ? m : n); \
|
|
int ldc = std::max<int>(1, m); \
|
|
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]; \
|
|
T* other_working_ptr = &other[i * other_stride]; \
|
|
PADDLE_ENFORCE_GPU_SUCCESS( \
|
|
dynload::rocsolver_##C##ormqr(handle, \
|
|
side, \
|
|
trans, \
|
|
m, \
|
|
n, \
|
|
k, \
|
|
a_working_ptr, \
|
|
lda, \
|
|
tau_working_ptr, \
|
|
other_working_ptr, \
|
|
ldc)); \
|
|
} \
|
|
}
|
|
FUNC_WITH_TYPES(ORMQR_BATCH_INSTANCE);
|
|
#endif
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
template <typename Context, typename T>
|
|
inline void BatchedOrmqr(const Context& dev_ctx,
|
|
bool left,
|
|
bool transpose,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
T* a,
|
|
int a_stride,
|
|
T* tau,
|
|
int tau_stride,
|
|
T* other,
|
|
int other_stride);
|
|
|
|
template <>
|
|
inline void BatchedOrmqr<GPUContext, float>(const GPUContext& dev_ctx,
|
|
bool left,
|
|
bool transpose,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
float* a,
|
|
int a_stride,
|
|
float* tau,
|
|
int tau_stride,
|
|
float* other,
|
|
int other_stride) {
|
|
int lwork = 0;
|
|
auto side = left ? CUBLAS_SIDE_LEFT : CUBLAS_SIDE_RIGHT;
|
|
auto trans = transpose ? CUBLAS_OP_T : CUBLAS_OP_N;
|
|
int lda = std::max<int>(1, left ? m : n);
|
|
int ldc = std::max<int>(1, m);
|
|
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnSormqr_bufferSize(
|
|
handle, side, trans, m, n, k, a, lda, tau, other, ldc, &lwork));
|
|
DenseTensor info;
|
|
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];
|
|
float* other_working_ptr = &other[i * other_stride];
|
|
|
|
handle = dev_ctx.cusolver_dn_handle();
|
|
DenseTensor workspace;
|
|
workspace.Resize({lwork});
|
|
float* workspace_ptr = dev_ctx.template Alloc<float>(&workspace);
|
|
|
|
// compute ormgr
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnSormqr(handle,
|
|
side,
|
|
trans,
|
|
m,
|
|
n,
|
|
k,
|
|
a_working_ptr,
|
|
lda,
|
|
tau_working_ptr,
|
|
other_working_ptr,
|
|
ldc,
|
|
workspace_ptr,
|
|
lwork,
|
|
info_d));
|
|
|
|
// 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 info is not zero but [%d]", i, info_h));
|
|
}
|
|
}
|
|
|
|
template <>
|
|
inline void BatchedOrmqr<GPUContext, double>(const GPUContext& dev_ctx,
|
|
bool left,
|
|
bool transpose,
|
|
int batch_size,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
double* a,
|
|
int a_stride,
|
|
double* tau,
|
|
int tau_stride,
|
|
double* other,
|
|
int other_stride) {
|
|
int lwork = 0;
|
|
auto side = left ? CUBLAS_SIDE_LEFT : CUBLAS_SIDE_RIGHT;
|
|
auto trans = transpose ? CUBLAS_OP_T : CUBLAS_OP_N;
|
|
int lda = std::max<int>(1, left ? m : n);
|
|
int ldc = std::max<int>(1, m);
|
|
|
|
auto handle = dev_ctx.cusolver_dn_handle();
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnDormqr_bufferSize(
|
|
handle, side, trans, m, n, k, a, lda, tau, other, ldc, &lwork));
|
|
DenseTensor info;
|
|
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];
|
|
double* other_working_ptr = &other[i * other_stride];
|
|
|
|
handle = dev_ctx.cusolver_dn_handle();
|
|
DenseTensor workspace;
|
|
workspace.Resize({lwork});
|
|
double* workspace_ptr = dev_ctx.template Alloc<double>(&workspace);
|
|
|
|
// compute ormgr
|
|
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnDormqr(handle,
|
|
side,
|
|
trans,
|
|
m,
|
|
n,
|
|
k,
|
|
a_working_ptr,
|
|
lda,
|
|
tau_working_ptr,
|
|
other_working_ptr,
|
|
ldc,
|
|
workspace_ptr,
|
|
lwork,
|
|
info_d));
|
|
|
|
// 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 info is not zero but [%d]", i, info_h));
|
|
}
|
|
}
|
|
#endif
|
|
|
|
} // namespace phi
|