wehub-resource-sync
462 lines
17 KiB
Plaintext
462 lines
17 KiB
Plaintext
// 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.
|
|
|
|
#include "paddle/phi/kernels/sparse/unary_kernel.h"
|
|
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
|
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
|
#include "paddle/phi/common/scalar.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/visit_type.h"
|
|
#include "paddle/phi/kernels/cast_kernel.h"
|
|
#include "paddle/phi/kernels/cum_kernel.h"
|
|
#include "paddle/phi/kernels/empty_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/elementwise_base.h"
|
|
#include "paddle/phi/kernels/index_select_kernel.h"
|
|
#include "paddle/phi/kernels/reduce_sum_kernel.h"
|
|
#include "paddle/phi/kernels/reshape_kernel.h"
|
|
#include "paddle/phi/kernels/sparse/empty_kernel.h"
|
|
#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
|
|
|
|
namespace phi {
|
|
namespace sparse {
|
|
|
|
template <typename T, typename IntT>
|
|
__global__ void SumCooCudaKernel(const IntT* x_indices_data,
|
|
const T* x_values_data,
|
|
const int64_t x_nnz,
|
|
const int64_t dense_dim,
|
|
const int64_t sparse_dim,
|
|
const int64_t axis,
|
|
const bool keep_dim,
|
|
IntT* out_indices_data,
|
|
T* out_values_data) {
|
|
CUDA_KERNEL_LOOP_TYPE(index_i, x_nnz, int64_t) {
|
|
int64_t i = 0;
|
|
for (int j = 0; j < dense_dim; ++j) {
|
|
out_values_data[j + index_i * dense_dim] = 0;
|
|
}
|
|
|
|
int64_t _index_j_ =
|
|
static_cast<int64_t>(blockIdx.y) * blockDim.y + threadIdx.y;
|
|
for (auto index_j = _index_j_; index_j < x_nnz;
|
|
index_j += static_cast<int64_t>(blockDim.y) * gridDim.y) {
|
|
// Determine whether the index_i and index_j elements have the same
|
|
// indices in all dimensions except for the specified axis dimension.
|
|
bool same = true;
|
|
for (int j = 0; j < sparse_dim + !keep_dim; ++j) {
|
|
if (j != axis && x_indices_data[index_i + j * x_nnz] !=
|
|
x_indices_data[index_j + j * x_nnz]) {
|
|
same = false;
|
|
break;
|
|
}
|
|
}
|
|
if (same) {
|
|
for (int j = 0; j < dense_dim; ++j) {
|
|
CudaAtomicAdd(&out_values_data[j + index_i * dense_dim],
|
|
x_values_data[j + index_j * dense_dim]);
|
|
}
|
|
}
|
|
}
|
|
if (_index_j_ != 0) {
|
|
return;
|
|
}
|
|
if (keep_dim) {
|
|
for (int j = 0; j < sparse_dim; ++j) {
|
|
if (j == axis) {
|
|
out_indices_data[index_i + j * x_nnz] = 0;
|
|
} else {
|
|
out_indices_data[index_i + j * x_nnz] =
|
|
x_indices_data[index_i + j * x_nnz];
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
for (int j = 0; j < sparse_dim; ++j) {
|
|
// out_indices_data [sparse_dim, x.nnz()]
|
|
int64_t x_indices_data_offset;
|
|
if (j < axis) {
|
|
x_indices_data_offset = index_i + j * x_nnz;
|
|
} else {
|
|
x_indices_data_offset = index_i + (j + 1) * x_nnz;
|
|
}
|
|
out_indices_data[index_i + j * x_nnz] =
|
|
x_indices_data[x_indices_data_offset];
|
|
}
|
|
}
|
|
}
|
|
|
|
__global__ void SumAllCsrCudaKernel(int64_t* out_crows_data,
|
|
int64_t* out_cols_data) {
|
|
CUDA_KERNEL_LOOP_TYPE(index, 2, int64_t) {
|
|
out_crows_data[index] = index;
|
|
if (index == 0) {
|
|
out_cols_data[0] = 0;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void SumCsr2DCudaKernel(const int64_t* x_crows_data,
|
|
const T* x_values_data,
|
|
const int64_t x_dim0,
|
|
int64_t* out_crows_data,
|
|
int64_t* out_cols_data,
|
|
T* out_values_data) {
|
|
CUDA_KERNEL_LOOP_TYPE(index, x_dim0 + 1, int64_t) {
|
|
out_crows_data[index] = index;
|
|
if (index != x_dim0) {
|
|
out_cols_data[index] = 0;
|
|
T sum_value = 0;
|
|
for (auto j = x_crows_data[index]; j < x_crows_data[index + 1]; ++j) {
|
|
sum_value += x_values_data[j];
|
|
}
|
|
out_values_data[index] = sum_value;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void SumCsr3DCudaKernel(const int64_t* x_crows_data,
|
|
const T* x_values_data,
|
|
const int64_t x_dim0,
|
|
const int64_t x_dim1,
|
|
const int64_t* batch_nnz_data,
|
|
int64_t* out_crows_data,
|
|
int64_t* out_cols_data,
|
|
T* out_values_data) {
|
|
{
|
|
CUDA_KERNEL_LOOP_TYPE(index, x_dim0 * x_dim1, int64_t) {
|
|
out_cols_data[index] = 0;
|
|
}
|
|
}
|
|
|
|
CUDA_KERNEL_LOOP_TYPE(index, x_dim0 * (x_dim1 + 1), int64_t) {
|
|
int64_t batch = index / (x_dim1 + 1);
|
|
int64_t number = index % (x_dim1 + 1);
|
|
out_crows_data[index] = number;
|
|
|
|
if (number != x_dim1) {
|
|
T sum_value = 0;
|
|
int64_t x_values_data_offset;
|
|
if (batch == 0) {
|
|
x_values_data_offset = 0;
|
|
} else {
|
|
x_values_data_offset = batch_nnz_data[batch - 1];
|
|
}
|
|
for (int64_t j = x_crows_data[index]; j < x_crows_data[index + 1]; ++j) {
|
|
sum_value += x_values_data[j + x_values_data_offset];
|
|
}
|
|
|
|
// `index - batch` would never exceed x_dim0 * x_dim1.
|
|
out_values_data[index - batch] = sum_value;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename IntT, typename Context>
|
|
void SumCooGPU0Kernel(const Context& dev_ctx,
|
|
const SparseCooTensor& x,
|
|
const IntArray& axis,
|
|
DataType dtype,
|
|
bool keep_dim,
|
|
SparseCooTensor* out) {
|
|
auto sparse_dim = x.sparse_dim();
|
|
// create out sparse tensor
|
|
const auto& x_dims = x.dims();
|
|
const auto& x_indices = x.indices();
|
|
const auto& x_values = x.values();
|
|
DDim out_dims;
|
|
DenseTensor out_indices;
|
|
DenseTensor out_values;
|
|
if (keep_dim) {
|
|
out_dims = make_ddim(std::vector<int64_t>(x_dims.size(), 1));
|
|
out_indices = Empty<IntT, Context>(dev_ctx, {sparse_dim, 1});
|
|
} else {
|
|
out_dims = make_ddim({1});
|
|
out_indices = Empty<IntT, Context>(dev_ctx, {1, 1});
|
|
}
|
|
funcs::SetConstant<Context, IntT> set_out_indices;
|
|
set_out_indices(dev_ctx, &out_indices, static_cast<IntT>(0));
|
|
out_values = phi::Sum<T>(dev_ctx, x.values(), {}, dtype, keep_dim);
|
|
out->SetMember(out_indices, out_values, out_dims, x.coalesced());
|
|
}
|
|
|
|
template <typename T, typename IntT, typename Context>
|
|
void SumCooGPU1Kernel(const Context& dev_ctx,
|
|
const SparseCooTensor& x,
|
|
const IntArray& axis,
|
|
DataType dtype,
|
|
bool keep_dim,
|
|
SparseCooTensor* out) {
|
|
auto sparse_dim = x.sparse_dim();
|
|
// create out sparse tensor
|
|
const auto& x_dims = x.dims();
|
|
const auto& x_indices = x.indices();
|
|
const auto& x_values = x.values();
|
|
DDim out_dims;
|
|
DenseTensor out_indices;
|
|
DenseTensor out_values;
|
|
auto n_dim = x.dims().size();
|
|
auto dim = axis[0] < 0 ? x_dims.size() + axis[0] : axis[0];
|
|
|
|
std::vector<int64_t> dims;
|
|
for (int i = 0; i < n_dim; ++i) {
|
|
if (i != dim) {
|
|
dims.emplace_back(x.dims()[i]);
|
|
} else if (keep_dim || (dim < sparse_dim && sparse_dim == 1)) {
|
|
dims.emplace_back(1);
|
|
}
|
|
}
|
|
out_dims = make_ddim(dims);
|
|
|
|
if (dim >= sparse_dim) {
|
|
out_indices = x_indices;
|
|
dim = dim - sparse_dim + 1;
|
|
out_values = phi::Sum<T>(dev_ctx, x.values(), {dim}, dtype, keep_dim);
|
|
out->SetMember(out_indices, out_values, out_dims, x.coalesced());
|
|
return;
|
|
}
|
|
|
|
// Ensure the sparse_dim is not less than 1.
|
|
if (sparse_dim == 1) {
|
|
keep_dim = true;
|
|
}
|
|
// if axis in sparse_dim and keep_dim, sparse_dim will be reduced.
|
|
if (!keep_dim) {
|
|
sparse_dim -= 1;
|
|
}
|
|
|
|
std::vector<int> out_values_dims;
|
|
out_values_dims.push_back(x.nnz());
|
|
for (auto i = 1; i < x.values().dims().size(); ++i) {
|
|
out_values_dims.push_back(static_cast<int>(x.values().dims()[i]));
|
|
}
|
|
int64_t dense_dim = std::accumulate(out_values_dims.begin() + 1,
|
|
out_values_dims.end(),
|
|
1,
|
|
std::multiplies<int64_t>());
|
|
|
|
out_indices = Empty<IntT, Context>(dev_ctx, {sparse_dim, x.nnz()});
|
|
out_values = Empty<T, Context>(dev_ctx, out_values_dims);
|
|
|
|
const auto* x_indices_data = x_indices.data<IntT>();
|
|
const auto* x_values_data = x_values.data<T>();
|
|
auto* out_indices_data = out_indices.data<IntT>();
|
|
auto* out_values_data = out_values.data<T>();
|
|
|
|
auto config = backends::gpu::GetGpuLaunchConfig2D(dev_ctx, x.nnz(), x.nnz());
|
|
SumCooCudaKernel<T, IntT><<<config.block_per_grid.x,
|
|
config.thread_per_block.x,
|
|
0,
|
|
dev_ctx.stream()>>>(x_indices_data,
|
|
x_values_data,
|
|
x.nnz(),
|
|
dense_dim,
|
|
sparse_dim,
|
|
dim,
|
|
keep_dim,
|
|
out_indices_data,
|
|
out_values_data);
|
|
if (dtype != phi::DataType::UNDEFINED && dtype != x.dtype()) {
|
|
out_values = Cast<T, Context>(dev_ctx, out_values, dtype);
|
|
}
|
|
out->SetMember(out_indices, out_values, out_dims, x.coalesced());
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void SumCooKernel(const Context& dev_ctx,
|
|
const SparseCooTensor& x,
|
|
const IntArray& axis,
|
|
DataType dtype,
|
|
bool keep_dim,
|
|
SparseCooTensor* out) {
|
|
const size_t n_dim = axis.size();
|
|
if (n_dim == 0) {
|
|
PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "SumCooGPUKernel", ([&] {
|
|
SumCooGPU0Kernel<T, data_t, Context>(
|
|
dev_ctx, x, axis, dtype, keep_dim, out);
|
|
}));
|
|
} else {
|
|
PD_VISIT_BASE_INTEGRAL_TYPES(x.indices().dtype(), "SumCooGPUKernel", ([&] {
|
|
SumCooGPU1Kernel<T, data_t, Context>(
|
|
dev_ctx, x, axis, dtype, keep_dim, out);
|
|
}));
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void SumCsr0Kernel(const Context& dev_ctx,
|
|
const SparseCsrTensor& x,
|
|
const IntArray& axis,
|
|
DataType dtype,
|
|
bool keep_dim,
|
|
SparseCsrTensor* out) {
|
|
auto x_dim0 = x.dims()[0];
|
|
auto x_dim1 = x.dims()[1];
|
|
const auto& x_crows = x.crows();
|
|
const auto& x_values = x.values();
|
|
const auto* x_crows_data = x_crows.data<int64_t>();
|
|
const auto* x_values_data = x_values.data<T>();
|
|
|
|
DenseTensor out_crows, out_cols, out_values;
|
|
DDim out_dims;
|
|
if (keep_dim && x.dims().size() == 3) {
|
|
out_dims = make_ddim({1, 1, 1});
|
|
} else {
|
|
out_dims = make_ddim({1, 1});
|
|
}
|
|
out_crows = Empty<int64_t, Context>(dev_ctx, {2}); // crows = [0, 1]
|
|
out_cols = Empty<int64_t, Context>(dev_ctx, {1}); // crows = [0]
|
|
auto* out_crows_data = out_crows.data<int64_t>();
|
|
auto* out_cols_data = out_cols.data<int64_t>();
|
|
|
|
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, 2, 1);
|
|
SumAllCsrCudaKernel<<<config.block_per_grid.x,
|
|
config.thread_per_block.x,
|
|
0,
|
|
dev_ctx.stream()>>>(out_crows_data, out_cols_data);
|
|
|
|
out_values = phi::Sum<T>(dev_ctx, x.values(), {}, dtype, true);
|
|
out->SetMember(out_crows, out_cols, out_values, out_dims);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void SumCsr1Kernel(const Context& dev_ctx,
|
|
const SparseCsrTensor& x,
|
|
const IntArray& axis,
|
|
DataType dtype,
|
|
bool keep_dim,
|
|
SparseCsrTensor* out) {
|
|
auto x_dim0 = x.dims()[0];
|
|
auto x_dim1 = x.dims()[1];
|
|
const auto& x_crows = x.crows();
|
|
const auto& x_values = x.values();
|
|
const auto* x_crows_data = x_crows.data<int64_t>();
|
|
const auto* x_values_data = x_values.data<T>();
|
|
|
|
DenseTensor out_crows, out_cols, out_values;
|
|
DDim out_dims;
|
|
out_crows = EmptyLike<int64_t, Context>(dev_ctx, x.crows());
|
|
auto* out_crows_data = out_crows.data<int64_t>();
|
|
|
|
if (x.dims().size() == 2) {
|
|
out_cols = Empty<int64_t, Context>(dev_ctx, {x_dim0});
|
|
out_values = Empty<T, Context>(dev_ctx, {x_dim0});
|
|
auto* out_cols_data = out_cols.data<int64_t>();
|
|
auto* out_values_data = out_values.data<T>();
|
|
out_dims = make_ddim({x_dim0, 1});
|
|
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x_dim0 + 1, 1);
|
|
SumCsr2DCudaKernel<T><<<config.block_per_grid.x,
|
|
config.thread_per_block.x,
|
|
0,
|
|
dev_ctx.stream()>>>(x_crows_data,
|
|
x_values_data,
|
|
x_dim0,
|
|
out_crows_data,
|
|
out_cols_data,
|
|
out_values_data);
|
|
|
|
} else {
|
|
out_cols = Empty<int64_t, Context>(dev_ctx, {x_dim0 * x_dim1});
|
|
out_values = Empty<T, Context>(dev_ctx, {x_dim0 * x_dim1});
|
|
auto* out_cols_data = out_cols.data<int64_t>();
|
|
auto* out_values_data = out_values.data<T>();
|
|
if (keep_dim) {
|
|
out_dims = make_ddim({x_dim0, x_dim1, 1});
|
|
} else {
|
|
out_dims = make_ddim({x_dim0, x_dim1});
|
|
}
|
|
|
|
DenseTensor x_crows_reshape =
|
|
Reshape<int64_t, Context>(dev_ctx, x_crows, {x_dim0, x_dim1 + 1});
|
|
DenseTensor last_indices = Empty<int64_t, Context>(dev_ctx, {1});
|
|
funcs::SetConstant<Context, int64_t> set_constant;
|
|
set_constant(dev_ctx, &last_indices, static_cast<int64_t>(x_dim1));
|
|
|
|
DenseTensor x_crows_last = Empty<int64_t, Context>(dev_ctx, {x_dim0, 1});
|
|
IndexSelectKernel<int64_t, Context>(
|
|
dev_ctx, x_crows_reshape, last_indices, 1, &x_crows_last);
|
|
|
|
DenseTensor batch_nnz = Empty<int64_t, Context>(dev_ctx, {x_dim0, 1});
|
|
CumsumKernel<int64_t, Context>(
|
|
dev_ctx, x_crows_last, Scalar(0), false, false, false, &batch_nnz);
|
|
auto* batch_nnz_data = batch_nnz.data<int64_t>();
|
|
|
|
auto config = backends::gpu::GetGpuLaunchConfig1D(
|
|
dev_ctx, x.dims()[0] * (x.dims()[1] + 1), 1);
|
|
SumCsr3DCudaKernel<T><<<config.block_per_grid.x,
|
|
config.thread_per_block.x,
|
|
0,
|
|
dev_ctx.stream()>>>(x_crows_data,
|
|
x_values_data,
|
|
x_dim0,
|
|
x_dim1,
|
|
batch_nnz_data,
|
|
out_crows_data,
|
|
out_cols_data,
|
|
out_values_data);
|
|
}
|
|
if (dtype != phi::DataType::UNDEFINED && dtype != x.dtype()) {
|
|
out_values = Cast<T, Context>(dev_ctx, out_values, dtype);
|
|
}
|
|
out->SetMember(out_crows, out_cols, out_values, out_dims);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void SumCsrKernel(const Context& dev_ctx,
|
|
const SparseCsrTensor& x,
|
|
const IntArray& axis,
|
|
DataType dtype,
|
|
bool keep_dim,
|
|
SparseCsrTensor* out) {
|
|
size_t n_dim = axis.size();
|
|
if (n_dim == 0) {
|
|
SumCsr0Kernel<T, Context>(dev_ctx, x, axis, dtype, keep_dim, out);
|
|
} else {
|
|
PADDLE_ENFORCE_EQ(axis[0],
|
|
-1,
|
|
common::errors::Unimplemented(
|
|
"`axis` of SumCsrKernel only support None or -1 now."
|
|
"More number will be supported in the future."));
|
|
SumCsr1Kernel<T, Context>(dev_ctx, x, axis, dtype, keep_dim, out);
|
|
}
|
|
}
|
|
|
|
} // namespace sparse
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(sum_coo,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::sparse::SumCooKernel,
|
|
float,
|
|
double,
|
|
int,
|
|
int64_t) {
|
|
kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
|
|
}
|
|
|
|
PD_REGISTER_KERNEL(sum_csr,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::sparse::SumCsrKernel,
|
|
float,
|
|
double,
|
|
int,
|
|
int64_t) {
|
|
kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
|
|
}
|