chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,233 @@
|
||||
// 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/backends/gpu/gpu_context.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/funcs/elementwise_base.h"
|
||||
#include "paddle/phi/kernels/reduce_sum_grad_kernel.h"
|
||||
#include "paddle/phi/kernels/reduce_sum_kernel.h"
|
||||
#include "paddle/phi/kernels/sparse/empty_kernel.h"
|
||||
#include "paddle/phi/kernels/sparse/unary_grad_kernel.h"
|
||||
#include "paddle/phi/kernels/sparse/unary_kernel.h"
|
||||
|
||||
namespace phi {
|
||||
namespace sparse {
|
||||
|
||||
template <typename T>
|
||||
__global__ void SetValueCudaKernel(const T* value,
|
||||
const int64_t length,
|
||||
T* data) {
|
||||
CUDA_KERNEL_LOOP_TYPE(index, length, int64_t) { data[index] = value[0]; }
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void SumCsr2DGradCudaKernel(const int64_t* x_crows_data,
|
||||
const T* dout_values_data,
|
||||
const int64_t x_dim0,
|
||||
T* dx_values_data) {
|
||||
// dout_crows_data[index] should be equal to index;
|
||||
CUDA_KERNEL_LOOP_TYPE(index, x_dim0, int64_t) {
|
||||
T value = dout_values_data[index];
|
||||
for (auto i = x_crows_data[index]; i < x_crows_data[index + 1]; ++i) {
|
||||
dx_values_data[i] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void SumCsr3DGradCudaKernel(const int64_t* x_crows_data,
|
||||
const T* dout_values_data,
|
||||
const int64_t x_dim0,
|
||||
const int64_t x_dim1,
|
||||
T* dx_values_data) {
|
||||
// dout_crows_data[index] should be equal to number;
|
||||
CUDA_KERNEL_LOOP_TYPE(index, x_dim0 * (x_dim1 + 1) - 1, int64_t) {
|
||||
int64_t batch = index / (x_dim1 + 1);
|
||||
int64_t number = index % (x_dim1 + 1);
|
||||
|
||||
// compute offset of dx_values_data in every batch
|
||||
int64_t batch_offset = 0;
|
||||
for (int64_t b = 1; b <= batch; ++b) {
|
||||
batch_offset += x_crows_data[b * (x_dim1 + 1) - 1];
|
||||
}
|
||||
|
||||
T value = dout_values_data[index - batch];
|
||||
for (auto i = x_crows_data[index]; i < x_crows_data[index + 1]; ++i) {
|
||||
dx_values_data[i + batch_offset] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IntT, typename Context>
|
||||
void SumCooGradGPUKernel(const Context& dev_ctx,
|
||||
const SparseCooTensor& x,
|
||||
const SparseCooTensor& dout,
|
||||
const IntArray& axis,
|
||||
bool keep_dim,
|
||||
SparseCooTensor* dx) {
|
||||
EmptyLikeCooKernel<T, Context>(dev_ctx, x, dx);
|
||||
unsigned int n_dim = axis.size();
|
||||
|
||||
const DenseTensor& x_indices = x.indices();
|
||||
const DenseTensor& dout_indices = dout.indices();
|
||||
const DenseTensor& dout_values = dout.values();
|
||||
const auto* dout_indices_data = dout_indices.data<IntT>();
|
||||
const auto* dout_values_data = dout_values.data<T>();
|
||||
|
||||
DenseTensor* dx_indices = dx->mutable_indices();
|
||||
DenseTensor* dx_values = dx->mutable_values();
|
||||
*dx_indices = x_indices;
|
||||
|
||||
const auto* dx_indices_data = dx_indices->data<IntT>();
|
||||
auto* dx_values_data = dx_values->data<T>();
|
||||
|
||||
if (n_dim == 0) {
|
||||
auto length = dx->nnz();
|
||||
for (auto i = 1; i < x.values().dims().size(); ++i) {
|
||||
length *= x.values().dims()[i];
|
||||
}
|
||||
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, length, 1);
|
||||
|
||||
SetValueCudaKernel<T>
|
||||
<<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
dev_ctx.stream()>>>(dout_values_data, length, dx_values_data);
|
||||
|
||||
if (dx_values->dtype() != dx->dtype()) {
|
||||
*dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype());
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
auto dim = axis[0] < 0 ? x.dims().size() + axis[0] : axis[0];
|
||||
auto sparse_dim = x.sparse_dim();
|
||||
if (dim >= sparse_dim) {
|
||||
dim = dim - sparse_dim + 1;
|
||||
phi::ReduceSumGradKernel<T, Context>(
|
||||
dev_ctx, x.values(), dout.values(), {dim}, keep_dim, false, dx_values);
|
||||
} else {
|
||||
*dx_values = dout_values;
|
||||
}
|
||||
if (dx_values->dtype() != dx->dtype()) {
|
||||
*dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype());
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SumCsrGradKernel(const Context& dev_ctx,
|
||||
const SparseCsrTensor& x,
|
||||
const SparseCsrTensor& dout,
|
||||
const IntArray& axis,
|
||||
bool keep_dim,
|
||||
SparseCsrTensor* dx) {
|
||||
EmptyLikeCsrKernel<T, Context>(dev_ctx, x, dx);
|
||||
size_t n_dim = axis.size();
|
||||
|
||||
const DenseTensor& x_crows = x.crows();
|
||||
const DenseTensor& x_cols = x.cols();
|
||||
const DenseTensor& dout_values = dout.values();
|
||||
|
||||
DenseTensor* dx_crows = dx->mutable_crows();
|
||||
DenseTensor* dx_cols = dx->mutable_cols();
|
||||
DenseTensor* dx_values = dx->mutable_values();
|
||||
|
||||
const auto* x_crows_data = x_crows.data<int64_t>();
|
||||
const auto* dout_values_data = dout_values.data<T>();
|
||||
auto* dx_values_data = dx_values->data<T>();
|
||||
|
||||
*dx_crows = x_crows;
|
||||
*dx_cols = x_cols;
|
||||
|
||||
if (n_dim == 0) {
|
||||
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, dx->nnz(), 1);
|
||||
SetValueCudaKernel<T>
|
||||
<<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
dev_ctx.stream()>>>(dout_values_data, dx->nnz(), dx_values_data);
|
||||
|
||||
if (dx_values->dtype() != dx->dtype()) {
|
||||
*dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype());
|
||||
}
|
||||
return;
|
||||
}
|
||||
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."));
|
||||
if (x.dims().size() == 2) {
|
||||
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x.dims()[0], 1);
|
||||
SumCsr2DGradCudaKernel<T><<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
dev_ctx.stream()>>>(
|
||||
x_crows_data, dout_values_data, x.dims()[0], dx_values_data);
|
||||
} else {
|
||||
auto config = backends::gpu::GetGpuLaunchConfig1D(
|
||||
dev_ctx, x.dims()[0] * (x.dims()[1] + 1), 1);
|
||||
SumCsr3DGradCudaKernel<T><<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
dev_ctx.stream()>>>(x_crows_data,
|
||||
dout_values_data,
|
||||
x.dims()[0],
|
||||
x.dims()[1],
|
||||
dx_values_data);
|
||||
}
|
||||
if (dx_values->dtype() != dx->dtype()) {
|
||||
*dx_values = Cast<T, Context>(dev_ctx, *dx_values, dx->dtype());
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SumCooGradKernel(const Context& dev_ctx,
|
||||
const SparseCooTensor& x,
|
||||
const SparseCooTensor& dout,
|
||||
const IntArray& axis,
|
||||
bool keep_dim,
|
||||
SparseCooTensor* dx) {
|
||||
PD_VISIT_BASE_INTEGRAL_TYPES(
|
||||
x.indices().dtype(), "SumCooGradGPUKernel", ([&] {
|
||||
SumCooGradGPUKernel<T, data_t, Context>(
|
||||
dev_ctx, x, dout, axis, keep_dim, dx);
|
||||
}));
|
||||
}
|
||||
} // namespace sparse
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(sum_coo_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::sparse::SumCooGradKernel,
|
||||
float,
|
||||
double,
|
||||
int16_t,
|
||||
int,
|
||||
int64_t,
|
||||
bool) {}
|
||||
|
||||
PD_REGISTER_KERNEL(sum_csr_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::sparse::SumCsrGradKernel,
|
||||
float,
|
||||
double,
|
||||
int16_t,
|
||||
int,
|
||||
int64_t,
|
||||
bool) {}
|
||||
Reference in New Issue
Block a user