wehub-resource-sync
406 lines
16 KiB
C++
406 lines
16 KiB
C++
// Copyright (c) 2023 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/ddim.h"
|
|
#include "paddle/phi/backends/dynload/rocsparse.h"
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/common/data_type.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/core/sparse_coo_tensor.h"
|
|
#include "paddle/phi/core/sparse_csr_tensor.h"
|
|
#include "paddle/phi/core/visit_type.h"
|
|
namespace phi {
|
|
namespace funcs {
|
|
namespace sparse {
|
|
|
|
template <typename IntT>
|
|
rocsparse_indextype GetGpuIndexType() {
|
|
if (std::is_same<IntT, int32_t>::value) {
|
|
return rocsparse_indextype_i32;
|
|
} else if (std::is_same<IntT, int64_t>::value) {
|
|
return rocsparse_indextype_i64;
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
rocsparse_datatype GetGpuDataType() {
|
|
if (std::is_same<T, float>::value) {
|
|
return rocsparse_datatype_f32_r;
|
|
} else if (std::is_same<T, double>::value) {
|
|
return rocsparse_datatype_f64_r;
|
|
}
|
|
}
|
|
|
|
inline rocsparse_operation GetTransposeOperation(const bool trans) {
|
|
if (trans) {
|
|
return rocsparse_operation_transpose;
|
|
} else {
|
|
return rocsparse_operation_none;
|
|
}
|
|
}
|
|
|
|
template <typename TensorType>
|
|
inline rocsparse_spmm_alg GetSpMMAlgorithm(const TensorType& x) {
|
|
return rocsparse_spmm_alg_default;
|
|
}
|
|
|
|
/************* SPARSE MATRIX DESCRIPTOR (COO/CSR) ************/
|
|
template <typename T, typename IntT>
|
|
inline void CreateCsrDescriptor(const SparseCsrTensor& x,
|
|
const GPUContext& dev_ctx,
|
|
rocsparse_spmat_descr* descriptor) {
|
|
std::vector<int64_t> xdim_vec = vectorize(x.dims());
|
|
auto x_ndims = xdim_vec.size();
|
|
PADDLE_ENFORCE_GE(
|
|
x_ndims,
|
|
2,
|
|
common::errors::InvalidArgument("the dim size of SparseCsrTensor must be "
|
|
"greater than or equal to 2."));
|
|
int64_t M = xdim_vec[x_ndims - 2];
|
|
int64_t N = xdim_vec[x_ndims - 1];
|
|
int batch_size = 1;
|
|
for (int i = 0; i < x_ndims - 2; i++) {
|
|
batch_size *= xdim_vec[i];
|
|
}
|
|
PADDLE_ENFORCE_EQ(x.non_zero_crows().numel(),
|
|
batch_size * (M + 1),
|
|
common::errors::PreconditionNotMet(
|
|
"the length of SparseCsrTensor crows is not right."));
|
|
|
|
const IntT* crows_data = x.non_zero_crows().data<IntT>();
|
|
const IntT* cols_data = x.non_zero_cols().data<IntT>();
|
|
const T* values_data = x.non_zero_elements().data<T>();
|
|
|
|
int64_t batch_nnz = x.nnz() / batch_size;
|
|
rocsparse_indextype itype = GetGpuIndexType<int64_t>();
|
|
rocsparse_indextype jtype = GetGpuIndexType<int64_t>();
|
|
rocsparse_datatype ttype = GetGpuDataType<T>();
|
|
dev_ctx.CusparseCall([&](rocsparse_handle handle) {
|
|
phi::dynload::rocsparse_create_csr_descr(descriptor,
|
|
M,
|
|
N,
|
|
batch_nnz,
|
|
const_cast<IntT*>(crows_data),
|
|
const_cast<IntT*>(cols_data),
|
|
const_cast<T*>(values_data),
|
|
itype,
|
|
jtype,
|
|
rocsparse_index_base_zero,
|
|
ttype);
|
|
});
|
|
if (batch_size > 1) {
|
|
// TODO(umiswing): Add batch sparse matmul support for ROCM after 5.2.0
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"Batch Sparse matmul use 'rocsparse_coo_set_strided_batch', which is "
|
|
"supported from ROCM 5.2.0"));
|
|
}
|
|
}
|
|
|
|
template <typename T, typename IntT>
|
|
inline void CreateCooDescriptor(const SparseCooTensor& x,
|
|
const GPUContext& dev_ctx,
|
|
rocsparse_spmat_descr* descriptor) {
|
|
std::vector<int64_t> xdim_vec = vectorize(x.dims());
|
|
auto x_ndims = xdim_vec.size();
|
|
PADDLE_ENFORCE_GE(
|
|
x_ndims,
|
|
2,
|
|
common::errors::InvalidArgument("the dim size of SparseCooTensor must be "
|
|
"greater than or equal to 2."));
|
|
|
|
int64_t M = xdim_vec[x_ndims - 2];
|
|
int64_t N = xdim_vec[x_ndims - 1];
|
|
int batch_size = 1;
|
|
for (int i = 0; i < x_ndims - 2; i++) {
|
|
batch_size *= xdim_vec[i];
|
|
}
|
|
int64_t nnz = x.nnz();
|
|
|
|
const IntT* indices_data = x.non_zero_indices().data<IntT>();
|
|
const T* values_data = x.non_zero_elements().data<T>();
|
|
auto rows_data = indices_data + (x_ndims - 2) * nnz;
|
|
auto cols_data = indices_data + (x_ndims - 1) * nnz;
|
|
|
|
int64_t batch_nnz = nnz / batch_size;
|
|
rocsparse_indextype itype = GetGpuIndexType<int64_t>();
|
|
rocsparse_datatype ttype = GetGpuDataType<T>();
|
|
dev_ctx.CusparseCall([&](rocsparse_handle handle) {
|
|
phi::dynload::rocsparse_create_coo_descr(descriptor,
|
|
M,
|
|
N,
|
|
batch_nnz,
|
|
const_cast<IntT*>(rows_data),
|
|
const_cast<IntT*>(cols_data),
|
|
const_cast<T*>(values_data),
|
|
itype,
|
|
rocsparse_index_base_zero,
|
|
ttype);
|
|
});
|
|
|
|
if (batch_size > 1) {
|
|
// TODO(umiswing): Add batch sparse matmul support for ROCM after 5.2.0
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"Batch Sparse matmul use 'rocsparse_coo_set_strided_batch', which is "
|
|
"supported from ROCM 5.2.0"));
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
class RocSparseSpMatDescriptor {
|
|
public:
|
|
explicit RocSparseSpMatDescriptor(const SparseCsrTensor& x,
|
|
const GPUContext& dev_ctx)
|
|
: dev_ctx_(dev_ctx) {
|
|
PD_VISIT_BASE_INTEGRAL_TYPES(
|
|
x.non_zero_crows().dtype(), "Csr RocSparseSpMatDescriptor", ([&] {
|
|
CreateCsrDescriptor<T, data_t>(x, dev_ctx_, &descriptor_);
|
|
}));
|
|
VLOG(6) << "Create csr rocsparse_spmat_descr " << &descriptor_;
|
|
}
|
|
explicit RocSparseSpMatDescriptor(const SparseCooTensor& x,
|
|
const GPUContext& dev_ctx)
|
|
: dev_ctx_(dev_ctx) {
|
|
PD_VISIT_BASE_INTEGRAL_TYPES(
|
|
x.non_zero_indices().dtype(), "Coo RocSparseSpMatDescriptor", ([&] {
|
|
CreateCooDescriptor<T, data_t>(x, dev_ctx_, &descriptor_);
|
|
}));
|
|
VLOG(6) << "Create coo rocsparse_spmat_descr " << &descriptor_;
|
|
}
|
|
|
|
~RocSparseSpMatDescriptor() {
|
|
dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
|
|
phi::dynload::rocsparse_destroy_spmat_descr(descriptor_);
|
|
});
|
|
VLOG(6) << "Destroy roscparse_spmat_descr " << &descriptor_;
|
|
}
|
|
|
|
const rocsparse_spmat_descr& descriptor() const { return descriptor_; }
|
|
|
|
private:
|
|
const GPUContext& dev_ctx_;
|
|
rocsparse_spmat_descr descriptor_;
|
|
};
|
|
|
|
/************* DENSE MATRIX DESCRIPTOR ************/
|
|
template <typename T>
|
|
class RocSparseDnMatDescriptor {
|
|
public:
|
|
explicit RocSparseDnMatDescriptor(const DenseTensor& x,
|
|
const GPUContext& dev_ctx)
|
|
: dev_ctx_(dev_ctx) {
|
|
std::vector<int64_t> xdim_vec = vectorize(x.dims());
|
|
auto x_ndims = xdim_vec.size();
|
|
PADDLE_ENFORCE_GE(
|
|
x_ndims,
|
|
2,
|
|
common::errors::InvalidArgument("the dim size of DenseTensor must be "
|
|
"greater than or equal to 2."));
|
|
|
|
int64_t M = xdim_vec[x_ndims - 2];
|
|
int64_t N = xdim_vec[x_ndims - 1];
|
|
int batch_size = 1;
|
|
for (int i = 0; i < x_ndims - 2; i++) {
|
|
batch_size *= xdim_vec[i];
|
|
}
|
|
|
|
const T* x_data = x.data<T>();
|
|
rocsparse_datatype ttype = GetGpuDataType<T>();
|
|
dev_ctx.CusparseCall([&](rocsparse_handle handle) {
|
|
phi::dynload::rocsparse_create_dnmat_descr(&descriptor_,
|
|
M,
|
|
N,
|
|
N,
|
|
const_cast<T*>(x_data),
|
|
ttype,
|
|
rocsparse_order_row);
|
|
});
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
x.numel(),
|
|
batch_size * M * N,
|
|
common::errors::InvalidArgument("The number of elements in DenseTensor "
|
|
"must equals to batch_size * M * N."));
|
|
if (batch_size > 1) {
|
|
// TODO(umiswing): Add batch sparse matmul support for ROCM after 5.2.0
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"Batch Sparse matmul use 'rocsparse_dnmat_set_strided_batch', which "
|
|
"is supported from ROCM 5.2.0"));
|
|
}
|
|
VLOG(6) << "Create cusparseDnMatDescr_t " << &descriptor_;
|
|
}
|
|
|
|
~RocSparseDnMatDescriptor() {
|
|
dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
|
|
phi::dynload::rocsparse_destroy_dnmat_descr(descriptor_);
|
|
});
|
|
VLOG(6) << "Destroy rocsparse_dnmat_descr " << &descriptor_;
|
|
}
|
|
|
|
const rocsparse_dnmat_descr& descriptor() const { return descriptor_; }
|
|
|
|
private:
|
|
const GPUContext& dev_ctx_;
|
|
rocsparse_dnmat_descr descriptor_;
|
|
};
|
|
|
|
/************* SPARSE*DENSE->DENSE MATMUL ************/
|
|
template <>
|
|
template <typename T, typename TensorType>
|
|
void SparseBlas<GPUContext>::SPMM(bool transa,
|
|
bool transb,
|
|
T alpha,
|
|
const TensorType& mat_a,
|
|
const DenseTensor& mat_b,
|
|
T beta,
|
|
DenseTensor* mat_out) const {
|
|
auto a_descriptor = RocSparseSpMatDescriptor<T>(mat_a, dev_ctx_);
|
|
auto b_descriptor = RocSparseDnMatDescriptor<T>(mat_b, dev_ctx_);
|
|
auto out_descriptor = RocSparseDnMatDescriptor<T>(*mat_out, dev_ctx_);
|
|
|
|
rocsparse_datatype ttype = GetGpuDataType<T>();
|
|
size_t buffer_size = 0;
|
|
|
|
// Query SpMM buffer
|
|
dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
|
|
phi::dynload::rocsparse_spmm(handle,
|
|
GetTransposeOperation(transa),
|
|
GetTransposeOperation(transb),
|
|
&alpha,
|
|
a_descriptor.descriptor(),
|
|
b_descriptor.descriptor(),
|
|
&beta,
|
|
out_descriptor.descriptor(),
|
|
ttype,
|
|
GetSpMMAlgorithm(mat_a),
|
|
rocsparse_spmm_stage_buffer_size,
|
|
&buffer_size,
|
|
nullptr);
|
|
});
|
|
|
|
// Allocate buffer
|
|
phi::Allocator::AllocationPtr tmp_buffer = phi::memory_utils::Alloc(
|
|
dev_ctx_.GetPlace(),
|
|
buffer_size,
|
|
phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx_.stream())));
|
|
void* tmp_buffer_ptr = tmp_buffer->ptr();
|
|
|
|
// Preprocess data
|
|
dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
|
|
phi::dynload::rocsparse_spmm(handle,
|
|
GetTransposeOperation(transa),
|
|
GetTransposeOperation(transb),
|
|
&alpha,
|
|
a_descriptor.descriptor(),
|
|
b_descriptor.descriptor(),
|
|
&beta,
|
|
out_descriptor.descriptor(),
|
|
ttype,
|
|
GetSpMMAlgorithm(mat_a),
|
|
rocsparse_spmm_stage_preprocess,
|
|
&buffer_size,
|
|
tmp_buffer_ptr);
|
|
});
|
|
|
|
// Performs the actual SpMM computation
|
|
dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
|
|
phi::dynload::rocsparse_spmm(handle,
|
|
GetTransposeOperation(transa),
|
|
GetTransposeOperation(transb),
|
|
&alpha,
|
|
a_descriptor.descriptor(),
|
|
b_descriptor.descriptor(),
|
|
&beta,
|
|
out_descriptor.descriptor(),
|
|
ttype,
|
|
GetSpMMAlgorithm(mat_a),
|
|
rocsparse_spmm_stage_compute,
|
|
&buffer_size,
|
|
tmp_buffer_ptr);
|
|
});
|
|
}
|
|
|
|
/************* DENSE*DENSE->SPARSE MATMUL ************/
|
|
#if HIP_VERSION >= 403
|
|
template <>
|
|
template <typename T, typename TensorType>
|
|
void SparseBlas<GPUContext>::SDDMM(bool transa,
|
|
bool transb,
|
|
T alpha,
|
|
const DenseTensor& mat_a,
|
|
const DenseTensor& mat_b,
|
|
T beta,
|
|
TensorType* mat_out) const {
|
|
auto a_descriptor = RocSparseDnMatDescriptor<T>(mat_a, dev_ctx_);
|
|
auto b_descriptor = RocSparseDnMatDescriptor<T>(mat_b, dev_ctx_);
|
|
auto out_descriptor = RocSparseSpMatDescriptor<T>(*mat_out, dev_ctx_);
|
|
|
|
rocsparse_datatype gpu_type = GetGpuDataType<T>();
|
|
size_t buffer_size = 0;
|
|
dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
|
|
phi::dynload::rocsparse_sddmm_buffer_size(handle,
|
|
GetTransposeOperation(transa),
|
|
GetTransposeOperation(transb),
|
|
&alpha,
|
|
a_descriptor.descriptor(),
|
|
b_descriptor.descriptor(),
|
|
&beta,
|
|
out_descriptor.descriptor(),
|
|
gpu_type,
|
|
rocsparse_sddmm_alg_default,
|
|
&buffer_size);
|
|
});
|
|
|
|
phi::Allocator::AllocationPtr tmp_buffer = phi::memory_utils::Alloc(
|
|
dev_ctx_.GetPlace(),
|
|
buffer_size,
|
|
phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx_.stream())));
|
|
void* tmp_buffer_ptr = tmp_buffer->ptr();
|
|
|
|
dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
|
|
phi::dynload::rocsparse_sddmm_preprocess(handle,
|
|
GetTransposeOperation(transa),
|
|
GetTransposeOperation(transb),
|
|
&alpha,
|
|
a_descriptor.descriptor(),
|
|
b_descriptor.descriptor(),
|
|
&beta,
|
|
out_descriptor.descriptor(),
|
|
gpu_type,
|
|
rocsparse_sddmm_alg_default,
|
|
tmp_buffer_ptr);
|
|
});
|
|
|
|
dev_ctx_.CusparseCall([&](rocsparse_handle handle) {
|
|
phi::dynload::rocsparse_sddmm(handle,
|
|
GetTransposeOperation(transa),
|
|
GetTransposeOperation(transb),
|
|
&alpha,
|
|
a_descriptor.descriptor(),
|
|
b_descriptor.descriptor(),
|
|
&beta,
|
|
out_descriptor.descriptor(),
|
|
gpu_type,
|
|
rocsparse_sddmm_alg_default,
|
|
tmp_buffer_ptr);
|
|
});
|
|
}
|
|
#endif
|
|
} // namespace sparse
|
|
} // namespace funcs
|
|
} // namespace phi
|